An open source codebase for object detection based on Jittor
import jittordet.datasets  # noqa: F401
import jittordet.engine  # noqa: F401
import jittordet.models  # noqa: F401
from .version import __version__

__all__ = ['__version__']

jittordet/__init__.py
JittorCV/jittordet
from .base import BaseDetDataset
from .coco import CocoDataset
from .samplers import (AspectRatioBatchSampler, BaseBatchSampler,
                       PadBatchSampler)
from .transforms import (LoadAnnotations, LoadImageFromFile, PackDetInputs,
                         RandomChoiceResize, RandomFlip, RandomResize, Resize)
from .voc import VocDataset
from .wrappers import ConcatDataset

__all__ = [
    'BaseDetDataset', 'CocoDataset', 'VocDataset', 'BaseBatchSampler',
    'PadBatchSampler', 'AspectRatioBatchSampler', 'PackDetInputs', 'Resize',
    'LoadAnnotations', 'LoadImageFromFile', 'RandomResize', 'RandomFlip',
    'RandomChoiceResize', 'ConcatDataset'
]

jittordet/datasets/__init__.py
JittorCV/jittordet
import copy
import os.path as osp

import jittor as jt
import numpy as np
from jittor.dataset import Dataset

from jittordet.engine import BATCH_SAMPLERS, TRANSFORMS
from jittordet.structures import DetDataSample, InstanceData


class Compose:
    """Modified from mmengine.dataset.base_dataest.

    Compose multiple transforms sequentially.
    """

    def __init__(self, transforms=None):
        if transforms is None:
            self.transforms = []
            return

        # validate data type
        if isinstance(transforms, dict):
            transforms = [transforms]

        self.transforms = []
        for transform in transforms:
            if isinstance(transform, dict):
                transform = TRANSFORMS.build(transform)
                if not callable(transform):
                    raise TypeError(f'transform should be a callable object, '
                                    f'but got {type(transform)}')
                self.transforms.append(transform)
            elif callable(transform):
                self.transforms.append(transform)
            else:
                raise TypeError(
                    f'transform must be a callable object or dict, '
                    f'but got {type(transform)}')

    def __call__(self, data):
        for t in self.transforms:
            data = t(data)
            if data is None:
                return None
        return data


class BaseDetDataset(Dataset):

    METAINFO = dict()

    def __init__(self,
                 batch_size,
                 num_workers=0,
                 data_root='',
                 data_path=None,
                 metainfo=None,
                 filter_cfg=None,
                 test_mode=False,
                 transforms=None,
                 batch_sampler=None,
                 max_refetch=100,
                 **kwargs):
        super().__init__(
            batch_size=batch_size, num_workers=num_workers, **kwargs)
        self.data_root = data_root
        self.filter_cfg = copy.deepcopy(filter_cfg)
        self.test_mode = test_mode
        self.max_refetch = max_refetch

        # load metainfo
        self._metainfo = self._load_metainfo(copy.deepcopy(metainfo))
        # datasetwise init
        self.init_datasetwise(data_path)
        # load data information
        self.data_list = self.load_data_list()
        # fliter illegal data
        self.data_list = self.filter_data()

        # set total length for jittor.utils.dataset
        self.total_len = len(self.data_list)

        # compose data transforms
        self.transforms = Compose(transforms)

        if batch_sampler is not None:
            self.batch_sampler = BATCH_SAMPLERS.build(
                batch_sampler, dataset=self)
        else:
            self.batch_sampler = None

    @property
    def metainfo(self):
        return copy.deepcopy(self._metainfo)

    @classmethod
    def _load_metainfo(cls, metainfo):
        cls_metainfo = copy.deepcopy(cls.METAINFO)
        if metainfo is None:
            return cls_metainfo
        if not isinstance(metainfo, dict):
            raise TypeError(
                f'metainfo should be a dict, but got {type(metainfo)}')

        cls_metainfo.update(metainfo)
        return cls_metainfo

    def init_datasetwise(self, datasetwise_cfg):
        if datasetwise_cfg is None:
            return
        assert isinstance(datasetwise_cfg, dict), \
            f'datasetwise_cfg must be a dict, but get {type(datasetwise_cfg)}'

        for k, v in datasetwise_cfg.items():
            if hasattr(self, k):
                raise RuntimeError(f'Attr {k} has been set in {type(self)}')

            def _join_data_root(value, data_root):
                if isinstance(value, list):
                    for i in range(len(value)):
                        value[i] = _join_data_root(value[i], data_root)
                elif isinstance(value, dict):
                    for key in value.keys():
                        value[key] = _join_data_root(value[key], data_root)
                elif isinstance(value, str):
                    if not osp.isabs(value):
                        value = osp.join(data_root, value)
                else:
                    raise TypeError(
                        'The contents in data_path should be str type.')
                return value

            v = _join_data_root(v, self.data_root)

            setattr(self, k, v)

    def load_data_list(self):
        raise NotImplementedError

    def filter_data(self):
        return self.data_list

    def prepare_data(self, idx):
        data_info = copy.deepcopy(self.data_list[idx])
        if idx >= 0:
            data_info['sample_idx'] = idx
        else:
            data_info['sample_idx'] = len(self.data_list) + idx

        data = self.transforms(data_info)
        return data

    def __getitem__(self, idx):
        if self.test_mode:
            data = self.prepare_data(idx)
            if data is None:
                raise Exception('Test time pipline should not get `None` '
                                'data_sample')
            return data

        for _ in range(self.max_refetch + 1):
            data = self.prepare_data(idx)
            # Broken images or random augmentations may cause the returned data
            # to be None
            if data is None:
                idx = np.random.randint(0, len(self.data_list))
                continue
            return data

        raise Exception(f'Cannot find valid image after {self.max_refetch}! '
                        'Please check your image path and pipeline')

    def __len__(self):
        if self.batch_sampler is not None:
            return len(self.batch_sampler)
        else:
            return super().__len__()

    def __batch_len__(self):
        if self.batch_sampler is not None:
            return self.batch_sampler.batch_len
        else:
            return super().__batch_len__()

    def _get_index_list(self):
        if self.batch_sampler is not None:
            index = self.batch_sampler.get_index_list(rng=self._shuffle_rng)
            self.real_len = len(index)
            self.batch_len = len(self.batch_sampler)
            world_size = 1 if not jt.in_mpi else jt.world_size
            self.real_batch_size = int(self.batch_size // world_size)
            return index
        else:
            return super()._get_index_list()

    def collate_batch(self, batch):
        """Override original `collate_batch` to disable stack."""
        if isinstance(batch[0], dict):
            new_batch = dict()
            for key in batch[0].keys():
                value = [data[key] for data in batch]
                value = self.collate_batch(value)
                new_batch[key] = value
        elif isinstance(batch[0], list):
            new_batch = list()
            for i in range(len(batch[0])):
                value = [data[i] for data in batch]
                value = self.collate_batch(value)
                new_batch.append(value)
        elif isinstance(batch[0], tuple):
            new_batch = list()
            for i in range(len(batch[0])):
                value = [data[i] for data in batch]
                value = self.collate_batch(value)
                new_batch.append(value)
            new_batch = tuple(new_batch)
        else:
            new_batch = batch
        return new_batch

    def to_jittor(self, batch):
        """Override to_jittor function in jittor.utils.dataset."""
        if self.keep_numpy_array:
            return batch
        if isinstance(batch, jt.Var):
            return batch
        if isinstance(batch, (InstanceData, DetDataSample)):
            return batch.to_jittor(self.stop_grad)
        if isinstance(batch, np.ndarray):
            batch = jt.array(batch)
            if self.stop_grad:
                batch = batch.stop_grad()
        if isinstance(batch, dict):
            return {k: self.to_jittor(v) for k, v in batch.items()}
        if isinstance(batch, list):
            return [self.to_jittor(v) for v in batch]
        if isinstance(batch, tuple):
            batch = [self.to_jittor(v) for v in batch]
            return tuple(batch)
        return batch

jittordet/datasets/base.py
JittorCV/jittordet
# Modified from mmdetection.dataset.coco
import copy
import os.path as osp

from pycocotools.coco import COCO

from ..engine import DATASETS
from .base import BaseDetDataset


@DATASETS.register_module()
class CocoDataset(BaseDetDataset):
    """Dataset for COCO."""

    METAINFO = {
        'classes':
        ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
         'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
         'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
         'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
         'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
         'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
         'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
         'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
         'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
         'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
         'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
         'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
         'scissors', 'teddy bear', 'hair drier', 'toothbrush'),
        # palette is a list of color tuples, which is used for visualization.
        'palette':
        [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),
         (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),
         (100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30),
         (165, 42, 42), (255, 77, 255), (0, 226, 252), (182, 182, 255),
         (0, 82, 0), (120, 166, 157), (110, 76, 0), (174, 57, 255),
         (199, 100, 0), (72, 0, 118), (255, 179, 240), (0, 125, 92),
         (209, 0, 151), (188, 208, 182), (0, 220, 176), (255, 99, 164),
         (92, 0, 73), (133, 129, 255), (78, 180, 255), (0, 228, 0),
         (174, 255, 243), (45, 89, 255), (134, 134, 103), (145, 148, 174),
         (255, 208, 186), (197, 226, 255), (171, 134, 1), (109, 63, 54),
         (207, 138, 255), (151, 0, 95), (9, 80, 61), (84, 105, 51),
         (74, 65, 105), (166, 196, 102), (208, 195, 210), (255, 109, 65),
         (0, 143, 149), (179, 0, 194), (209, 99, 106), (5, 121, 0),
         (227, 255, 205), (147, 186, 208), (153, 69, 1), (3, 95, 161),
         (163, 255, 0), (119, 0, 170), (0, 182, 199), (0, 165, 120),
         (183, 130, 88), (95, 32, 0), (130, 114, 135), (110, 129, 133),
         (166, 74, 118), (219, 142, 185), (79, 210, 114), (178, 90, 62),
         (65, 70, 15), (127, 167, 115), (59, 105, 106), (142, 108, 45),
         (196, 172, 0), (95, 54, 80), (128, 76, 255), (201, 57, 1),
         (246, 0, 122), (191, 162, 208)]
    }

    def load_data_list(self):
        self.coco = COCO(self.ann_file)
        # The order of returned `cat_ids` will not
        # change with the order of the `classes`
        self.cat_ids = self.coco.getCatIds(self.metainfo['classes'])
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}
        self.cat_img_map = copy.deepcopy(self.coco.catToImgs)

        img_ids = self.coco.getImgIds()
        data_list = []
        for img_id in img_ids:
            raw_img_info = self.coco.loadImgs([img_id])[0]
            raw_img_info['img_id'] = img_id

            ann_ids = self.coco.getAnnIds(imgIds=[img_id])
            raw_ann_info = self.coco.loadAnns(ann_ids)

            parsed_data_info = self.parse_data_info(raw_ann_info, raw_img_info)
            data_list.append(parsed_data_info)

        del self.coco

        return data_list

    def parse_data_info(self, ann_info, img_info):
        data_info = {}

        img_path = osp.join(self.img_path, img_info['file_name'])
        data_info['img_path'] = img_path
        data_info['img_id'] = img_info['img_id']
        data_info['height'] = img_info['height']
        data_info['width'] = img_info['width']

        instances = []
        for i, ann in enumerate(ann_info):
            instance = {}

            if ann.get('ignore', False):
                continue
            x1, y1, w, h = ann['bbox']
            inter_w = max(0, min(x1 + w, img_info['width']) - max(x1, 0))
            inter_h = max(0, min(y1 + h, img_info['height']) - max(y1, 0))
            if inter_w * inter_h == 0:
                continue
            if ann['area'] <= 0 or w < 1 or h < 1:
                continue
            if ann['category_id'] not in self.cat_ids:
                continue
            bbox = [x1, y1, x1 + w, y1 + h]

            if ann.get('iscrowd', False):
                instance['ignore_flag'] = 1
            else:
                instance['ignore_flag'] = 0
            instance['bbox'] = bbox
            instance['bbox_label'] = self.cat2label[ann['category_id']]

            if ann.get('segmentation', None):
                instance['mask'] = ann['segmentation']

            instances.append(instance)
        data_info['instances'] = instances
        return data_info

    def filter_data(self):
        if self.test_mode:
            return self.data_list

        if self.filter_cfg is None:
            return self.data_list

        filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False)
        min_size = self.filter_cfg.get('min_size', 0)

        # obtain images that contain annotation
        ids_with_ann = set(data_info['img_id'] for data_info in self.data_list)
        # obtain images that contain annotations of the required categories
        ids_in_cat = set()
        for i, class_id in enumerate(self.cat_ids):
            ids_in_cat |= set(self.cat_img_map[class_id])
        # merge the image id sets of the two conditions and use the merged set
        # to filter out images if self.filter_empty_gt=True
        ids_in_cat &= ids_with_ann

        valid_data_infos = []
        for i, data_info in enumerate(self.data_list):
            img_id = data_info['img_id']
            width = data_info['width']
            height = data_info['height']
            if filter_empty_gt and img_id not in ids_in_cat:
                continue
            if min(width, height) >= min_size:
                valid_data_infos.append(data_info)

        return valid_data_infos

jittordet/datasets/coco.py
JittorCV/jittordet
from .aspect_ratio_batch_sampler import AspectRatioBatchSampler
from .base_batch_sampler import BaseBatchSampler
from .pad_batch_sampler import PadBatchSampler

__all__ = ['BaseBatchSampler', 'PadBatchSampler', 'AspectRatioBatchSampler']

jittordet/datasets/samplers/__init__.py
JittorCV/jittordet
import jittor as jt
import numpy as np

from jittordet.engine import BATCH_SAMPLERS
from .base_batch_sampler import BaseBatchSampler


@BATCH_SAMPLERS.register_module()
class AspectRatioBatchSampler(BaseBatchSampler):

    def __init__(self, dataset, drop_last=False):
        super().__init__(dataset=dataset)
        self.drop_last = drop_last
        # statstic different aspect ratios idx.
        idx_bucket1, idx_bucket2 = [], []
        for idx, data in enumerate(dataset.data_list):
            if data['width'] > data['height']:
                idx_bucket1.append(idx)
            else:
                idx_bucket2.append(idx)
        self.idx_bucket1 = np.array(idx_bucket1)
        self.idx_bucket2 = np.array(idx_bucket2)

    def get_index_list(self, rng=None):
        if rng is None:
            rng = np.random.default_rng()

        # shuffle
        shuffle_idx = rng.permutation(self.idx_bucket1.size)
        idx_bucket1 = self.idx_bucket1[shuffle_idx]
        shuffle_idx = rng.permutation(self.idx_bucket2.size)
        idx_bucket2 = self.idx_bucket2[shuffle_idx]

        # drop last size
        world_size = 1 if not jt.in_mpi else jt.world_size
        total_bs = self.total_bs
        real_bs = int(total_bs // world_size)
        assert real_bs * jt.world_size == total_bs
        if idx_bucket1.size % real_bs != 0:
            mod_size = idx_bucket1.size % real_bs
            idx_bucket1 = idx_bucket1[:-mod_size]
        idx_bucket1 = idx_bucket1.reshape(-1, real_bs)
        if idx_bucket2.size % real_bs != 0:
            mod_size = idx_bucket2.size % real_bs
            idx_bucket2 = idx_bucket2[:-mod_size]
        idx_bucket2 = idx_bucket2.reshape(-1, real_bs)

        index = np.concatenate([idx_bucket1, idx_bucket2], axis=0)
        shuffle_idx = rng.permutation(index.shape[0])
        index = index[shuffle_idx]

        real_bs_num = len(self) * world_size
        repeat_num = int((real_bs_num - 0.5) // index.shape[0] + 1)
        index = np.concatenate([index] * repeat_num, axis=0)
        index = index[:real_bs_num]

        if jt.in_mpi:
            index = index.reshape(-1, total_bs)
            index = index[:, jt.rank * real_bs:(jt.rank + 1) * real_bs]

        index = index.flatten()
        return index

jittordet/datasets/samplers/aspect_ratio_batch_sampler.py
JittorCV/jittordet
from abc import ABCMeta, abstractmethod

from jittordet.engine import BATCH_SAMPLERS


@BATCH_SAMPLERS.register_module()
class BaseBatchSampler(metaclass=ABCMeta):

    def __init__(self, dataset):
        self.total_bs = dataset.batch_size
        self.num_data_list = len(dataset.data_list)

    def __len__(self):
        length = int((self.num_data_list - 0.5) // self.total_bs)
        if hasattr(self, 'drop_last') and not self.drop_last:
            length += 1
        return length

    @abstractmethod
    def get_index_list(self, rng=None):
        pass

jittordet/datasets/samplers/base_batch_sampler.py
JittorCV/jittordet
from math import ceil

import jittor as jt
import numpy as np

from jittordet.engine import BATCH_SAMPLERS
from .base_batch_sampler import BaseBatchSampler


@BATCH_SAMPLERS.register_module()
class PadBatchSampler(BaseBatchSampler):

    def __init__(self, dataset, shuffle=False, drop_last=False):
        super().__init__(dataset=dataset)
        self.shuffle = shuffle
        self.drop_last = drop_last

    def get_index_list(self, rng=None):
        if rng is None:
            rng = np.random.default_rng()

        index = rng.permutation(self.num_data_list) if self.shuffle \
            else np.arange(self.num_data_list)

        mod_size = self.num_data_list % self.total_bs
        if mod_size != 0:
            if self.drop_last:
                index = index[:-mod_size]
            else:
                padded_size = int(
                    ceil(self.num_data_list / self.total_bs) * self.total_bs)
                # repeat index to avoid data_list is shorter than batch size
                repeat_num = int((self.total_bs - 0.5) // self.num_data_list +
                                 1)
                repeat_num = max(repeat_num, 2)
                index = np.concatenate([index] * repeat_num)
                index = index[:padded_size]

        if jt.in_mpi:
            rank, world_size = jt.rank, jt.world_size
            real_bs = int(self.total_bs // world_size)
            assert real_bs * world_size == self.total_bs
            index = index.reshape(-1, self.total_bs)
            index = index[:, rank * real_bs:(rank + 1) * real_bs]
            index = index.flatten()

        return index

jittordet/datasets/samplers/pad_batch_sampler.py
JittorCV/jittordet
from .formatting import PackDetInputs
from .loading import LoadAnnotations, LoadImageFromFile
from .transforms import RandomChoiceResize, RandomFlip, RandomResize, Resize

__all__ = [
    'PackDetInputs', 'LoadAnnotations', 'LoadImageFromFile', 'Resize',
    'RandomResize', 'RandomChoiceResize', 'RandomFlip'
]

jittordet/datasets/transforms/__init__.py
JittorCV/jittordet
import numpy as np

from jittordet.engine import TRANSFORMS
from jittordet.structures import DetDataSample, InstanceData


@TRANSFORMS.register_module()
class PackDetInputs:
    """Pack the inputs data for detection."""
    mapping_table = {
        'gt_bboxes': 'bboxes',
        'gt_bboxes_labels': 'labels',
    }

    def __init__(self,
                 meta_keys=('sample_idx', 'img_id', 'img_path', 'ori_shape',
                            'img_shape', 'scale_factor', 'flip',
                            'flip_direction')):
        self.meta_keys = meta_keys

    def __call__(self, results):
        packed_results = dict()
        if 'img' in results:
            img = results['img']
            if len(img.shape) < 3:
                img = np.expand_dims(img, -1)
            img = np.ascontiguousarray(img.transpose(2, 0, 1))
            packed_results['inputs'] = img

        if 'gt_ignore_flags' in results:
            valid_idx = np.where(results['gt_ignore_flags'] == 0)[0]
            ignore_idx = np.where(results['gt_ignore_flags'] == 1)[0]

        data_sample = DetDataSample()
        instance_data = InstanceData()
        ignore_instance_data = InstanceData()

        for key in self.mapping_table.keys():
            if key not in results:
                continue
            if 'gt_ignore_flags' in results:
                instance_data[
                    self.mapping_table[key]] = results[key][valid_idx]
                ignore_instance_data[
                    self.mapping_table[key]] = results[key][ignore_idx]
            else:
                instance_data[self.mapping_table[key]] = results[key]
        data_sample.gt_instances = instance_data
        data_sample.ignored_instances = ignore_instance_data

        img_meta = {}
        for key in self.meta_keys:
            assert key in results, f'`{key}` is not found in `results`, ' \
                f'the valid keys are {list(results)}.'
            img_meta[key] = results[key]

        data_sample.set_metainfo(img_meta)
        packed_results['data_samples'] = data_sample

        return packed_results

    def __repr__(self) -> str:
        repr_str = self.__class__.__name__
        repr_str += f'(meta_keys={self.meta_keys})'
        return repr_str

jittordet/datasets/transforms/formatting.py
JittorCV/jittordet
import cv2
import numpy as np

from jittordet.engine import TRANSFORMS


@TRANSFORMS.register_module()
class LoadImageFromFile:
    """Load an image from file."""

    def __init__(self, to_float32=False):
        self.to_float32 = to_float32

    def __call__(self, results):
        img = cv2.imread(results['img_path'])
        if self.to_float32:
            img = img.astype(np.float32)

        results['img'] = img
        results['img_shape'] = img.shape[:2]
        results['ori_shape'] = img.shape[:2]
        return results

    def __repr__(self):
        repr_str = (f'{self.__class__.__name__}('
                    f'to_float32={self.to_float32})')
        return repr_str


@TRANSFORMS.register_module()
class LoadAnnotations:
    """Load multiple types of annotations."""

    def __init__(self, with_bbox=True, with_label=True):
        self.with_bbox = with_bbox
        self.with_label = with_label

    def _load_bboxes(self, results):
        gt_bboxes = []
        gt_ignore_flags = []
        for instance in results.get('instances', []):
            gt_bboxes.append(instance['bbox'])
            gt_ignore_flags.append(instance['ignore_flag'])

        results['gt_bboxes'] = np.array(
            gt_bboxes, dtype=np.float32).reshape((-1, 4))
        results['gt_ignore_flags'] = np.array(gt_ignore_flags, dtype=np.bool)

    def _load_labels(self, results):
        gt_bboxes_labels = []
        for instance in results.get('instances', []):
            gt_bboxes_labels.append(instance['bbox_label'])
        results['gt_bboxes_labels'] = np.array(
            gt_bboxes_labels, dtype=np.int64)

    def __call__(self, results):
        if self.with_bbox:
            self._load_bboxes(results)
        if self.with_label:
            self._load_labels(results)
        return results

    def __repr__(self) -> str:
        repr_str = self.__class__.__name__
        repr_str += f'(with_bbox={self.with_bbox}, '
        repr_str += f'with_label={self.with_label})'
        return repr_str

jittordet/datasets/transforms/loading.py
JittorCV/jittordet
from collections.abc import Sequence

import numpy as np

from jittordet.engine import TRANSFORMS
from jittordet.utils import (_scale_size, imflip, imrescale, imresize,
                             is_list_of, is_seq_of, is_tuple_of)


@TRANSFORMS.register_module()
class Resize:

    def __init__(self,
                 scale,
                 scale_factor=None,
                 keep_ratio=False,
                 clip_object_border=True,
                 backend='cv2',
                 interpolation='bilinear'):
        assert scale is not None or scale_factor is not None, (
            '`scale` and'
            '`scale_factor` can not both be `None`')
        if scale is None:
            self.scale = None
        else:
            if isinstance(scale, int):
                self.scale = (scale, scale)
            else:
                self.scale = scale

        self.backend = backend
        self.interpolation = interpolation
        self.keep_ratio = keep_ratio
        self.clip_object_border = clip_object_border
        if scale_factor is None:
            self.scale_factor = None
        elif isinstance(scale_factor, float):
            self.scale_factor = (scale_factor, scale_factor)
        elif isinstance(scale_factor, tuple):
            assert (len(scale_factor)) == 2
            self.scale_factor = scale_factor
        else:
            raise TypeError(
                f'expect scale_factor is float or Tuple(float), but'
                f'get {type(scale_factor)}')

    def _resize_img(self, results: dict) -> None:
        """Resize images with ``results['scale']``."""

        if results.get('img', None) is not None:
            if self.keep_ratio:
                img, scale_factor = imrescale(
                    results['img'],
                    results['scale'],
                    interpolation=self.interpolation,
                    return_scale=True,
                    backend=self.backend)
                # the w_scale and h_scale has minor difference
                # a real fix should be done in the mmcv.imrescale in the future
                new_h, new_w = img.shape[:2]
                h, w = results['img'].shape[:2]
                w_scale = new_w / w
                h_scale = new_h / h
            else:
                img, w_scale, h_scale = imresize(
                    results['img'],
                    results['scale'],
                    interpolation=self.interpolation,
                    return_scale=True,
                    backend=self.backend)
            results['img'] = img
            results['img_shape'] = img.shape[:2]
            results['scale_factor'] = (w_scale, h_scale)
            results['keep_ratio'] = self.keep_ratio

    def _resize_bboxes(self, results: dict) -> None:
        """Resize bounding boxes with ``results['scale_factor']``."""
        if results.get('gt_bboxes', None) is not None:
            bboxes = results['gt_bboxes'] * np.tile(
                np.array(results['scale_factor']), 2)
            if self.clip_object_border:
                bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0,
                                          results['img_shape'][1])
                bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0,
                                          results['img_shape'][0])
            results['gt_bboxes'] = bboxes

    def __call__(self, results):
        if self.scale:
            results['scale'] = self.scale
        else:
            img_shape = results['img'].shape[:2]
            results['scale'] = _scale_size(img_shape[::-1], self.scale_factor)
        self._resize_img(results)
        self._resize_bboxes(results)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(scale={self.scale}, '
        repr_str += f'scale_factor={self.scale_factor}, '
        repr_str += f'keep_ratio={self.keep_ratio}, '
        repr_str += f'clip_object_border={self.clip_object_border}), '
        repr_str += f'backend={self.backend}), '
        repr_str += f'interpolation={self.interpolation})'
        return repr_str


@TRANSFORMS.register_module()
class RandomResize:
    """Random resize images & bboxs."""

    def __init__(self,
                 scale,
                 ratio_range=None,
                 resize_type='Resize',
                 **resize_kwargs):

        self.scale = scale
        self.ratio_range = ratio_range

        self.resize_cfg = dict(type=resize_type, **resize_kwargs)
        # create a empty Reisize object
        self.resize = TRANSFORMS.build({'scale': 0, **self.resize_cfg})

    @staticmethod
    def _random_sample(scales):
        """Randomly sample a scale from a list of tuples."""
        assert is_list_of(scales, tuple) and len(scales) == 2
        scale_0 = [scales[0][0], scales[1][0]]
        scale_1 = [scales[0][1], scales[1][1]]
        edge_0 = np.random.randint(min(scale_0), max(scale_0) + 1)
        edge_1 = np.random.randint(min(scale_1), max(scale_1) + 1)
        scale = (edge_0, edge_1)
        return scale

    @staticmethod
    def _random_sample_ratio(scale, ratio_range):
        """Randomly sample a scale from a tuple."""
        assert isinstance(scale, tuple) and len(scale) == 2
        min_ratio, max_ratio = ratio_range
        assert min_ratio <= max_ratio
        ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio
        scale = int(scale[0] * ratio), int(scale[1] * ratio)
        return scale

    def _random_scale(self) -> tuple:
        """Randomly sample an scale according to the type of ``scale``."""
        if is_tuple_of(self.scale, int):
            assert self.ratio_range is not None and len(self.ratio_range) == 2
            scale = self._random_sample_ratio(
                self.scale,  # type: ignore
                self.ratio_range)
        elif is_seq_of(self.scale, tuple):
            scale = self._random_sample(self.scale)  # type: ignore
        else:
            raise NotImplementedError('Do not support sampling function '
                                      f'for "{self.scale}"')

        return scale

    def __call__(self, results):
        results['scale'] = self._random_scale()
        self.resize.scale = results['scale']
        results = self.resize(results)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(scale={self.scale}, '
        repr_str += f'ratio_range={self.ratio_range}, '
        repr_str += f'resize_cfg={self.resize_cfg})'
        return repr_str


@TRANSFORMS.register_module()
class RandomChoiceResize:
    """Resize images & bbox from a list of multiple scales."""

    def __init__(self, scales, resize_type='Resize', **resize_kwargs):
        if isinstance(scales, list):
            self.scales = scales
        else:
            self.scales = [scales]
        assert is_list_of(self.scales, tuple)

        self.resize_cfg = dict(type=resize_type, **resize_kwargs)
        # create a empty Resize object
        self.resize = TRANSFORMS.build({'scale': 0, **self.resize_cfg})

    def _random_select(self):
        """Randomly select an scale from given candidates."""
        assert is_list_of(self.scales, tuple)
        scale_idx = np.random.randint(len(self.scales))
        scale = self.scales[scale_idx]
        return scale, scale_idx

    def __call__(self, results):
        """Apply resize transforms on results from a list of scales."""
        target_scale, scale_idx = self._random_select()
        self.resize.scale = target_scale
        results = self.resize(results)
        results['scale_idx'] = scale_idx
        return results

    def __repr__(self) -> str:
        repr_str = self.__class__.__name__
        repr_str += f'(scales={self.scales}'
        repr_str += f', resize_cfg={self.resize_cfg})'
        return repr_str


@TRANSFORMS.register_module()
class RandomFlip:
    """Flip the image & bbox."""

    def __init__(self, prob=None, direction='horizontal'):
        if isinstance(prob, list):
            assert is_list_of(prob, float)
            assert 0 <= sum(prob) <= 1
        elif isinstance(prob, float):
            assert 0 <= prob <= 1
        else:
            raise ValueError(f'probs must be float or list of float, but \
                              got `{type(prob)}`.')
        self.prob = prob

        valid_directions = ['horizontal', 'vertical', 'diagonal']
        if isinstance(direction, str):
            assert direction in valid_directions
        elif isinstance(direction, list):
            assert is_list_of(direction, str)
            assert set(direction).issubset(set(valid_directions))
        else:
            raise ValueError(f'direction must be either str or list of str, \
                               but got `{type(direction)}`.')
        self.direction = direction

        if isinstance(prob, list):
            assert len(prob) == len(self.direction)

    def _flip_bbox(self, bboxes: np.ndarray, img_shape, direction: str):
        """Flip bboxes horizontally."""
        assert bboxes.shape[-1] % 4 == 0
        flipped = bboxes.copy()
        h, w = img_shape
        if direction == 'horizontal':
            flipped[..., 0::4] = w - bboxes[..., 2::4]
            flipped[..., 2::4] = w - bboxes[..., 0::4]
        elif direction == 'vertical':
            flipped[..., 1::4] = h - bboxes[..., 3::4]
            flipped[..., 3::4] = h - bboxes[..., 1::4]
        elif direction == 'diagonal':
            flipped[..., 0::4] = w - bboxes[..., 2::4]
            flipped[..., 1::4] = h - bboxes[..., 3::4]
            flipped[..., 2::4] = w - bboxes[..., 0::4]
            flipped[..., 3::4] = h - bboxes[..., 1::4]
        else:
            raise ValueError(
                f"Flipping direction must be 'horizontal', 'vertical', \
                  or 'diagonal', but got '{direction}'")
        return flipped

    def _choose_direction(self) -> str:
        """Choose the flip direction according to `prob` and `direction`"""
        if isinstance(self.direction,
                      Sequence) and not isinstance(self.direction, str):
            # None means non-flip
            direction_list: list = list(self.direction) + [None]
        elif isinstance(self.direction, str):
            # None means non-flip
            direction_list = [self.direction, None]

        if isinstance(self.prob, list):
            non_prob: float = 1 - sum(self.prob)
            prob_list = self.prob + [non_prob]
        elif isinstance(self.prob, float):
            non_prob = 1. - self.prob
            # exclude non-flip
            single_ratio = self.prob / (len(direction_list) - 1)
            prob_list = [single_ratio] * (len(direction_list) - 1) + [non_prob]

        cur_dir = np.random.choice(direction_list, p=prob_list)

        return cur_dir

    def _flip(self, results: dict) -> None:
        """Flip images, bounding boxes, semantic segmentation map and
        keypoints."""
        # flip image
        results['img'] = imflip(
            results['img'], direction=results['flip_direction'])

        img_shape = results['img'].shape[:2]

        # flip bboxes
        if results.get('gt_bboxes', None) is not None:
            results['gt_bboxes'] = self._flip_bbox(results['gt_bboxes'],
                                                   img_shape,
                                                   results['flip_direction'])

    def __call__(self, results):
        cur_dir = self._choose_direction()
        if cur_dir is None:
            results['flip'] = False
            results['flip_direction'] = None
        else:
            results['flip'] = True
            results['flip_direction'] = cur_dir
            self._flip(results)
        return results

    def __repr__(self) -> str:
        repr_str = self.__class__.__name__
        repr_str += f'(prob={self.prob}, '
        repr_str += f'direction={self.direction})'
        return repr_str

jittordet/datasets/transforms/transforms.py
JittorCV/jittordet
# modified from mmdetection.datasets.xml_style
import os.path as osp
import xml.etree.ElementTree as ET

from PIL import Image

from ..engine import DATASETS
from .base import BaseDetDataset


@DATASETS.register_module()
class VocDataset(BaseDetDataset):

    METAINFO = {
        'classes':
        ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat',
         'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person',
         'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'),
        # PALETTE is a list of color tuples, which is used for visualization.
        'palette': [(106, 0, 228), (119, 11, 32), (165, 42, 42), (0, 0, 192),
                    (197, 226, 255), (0, 60, 100), (0, 0, 142), (255, 77, 255),
                    (153, 69, 1), (120, 166, 157), (0, 182, 199),
                    (0, 226, 252), (182, 182, 255), (0, 0, 230), (220, 20, 60),
                    (163, 255, 0), (0, 82, 0), (3, 95, 161), (0, 80, 100),
                    (183, 130, 88)]
    }

    def load_data_list(self):
        assert self.metainfo.get('classes', None) is not None, \
            'CLASSES in `VocDataset` can not be None.'
        self.cat2label = {
            cat: i
            for i, cat in enumerate(self.metainfo['classes'])
        }

        data_list = []
        for img_id in open(self.ann_file, 'r'):
            img_id = img_id.strip()
            img_path = osp.join(self.img_path, f'{img_id}.jpg')
            xml_path = osp.join(self.xml_path, f'{img_id}.xml')

            raw_img_info = {}
            raw_img_info['img_id'] = img_id
            raw_img_info['img_path'] = img_path
            raw_img_info['xml_path'] = xml_path

            parsed_data_info = self.parse_data_info(raw_img_info)
            data_list.append(parsed_data_info)
        return data_list

    def parse_data_info(self, img_info):
        data_info = {}
        data_info['img_id'] = img_info['img_id']
        data_info['img_path'] = img_info['img_path']
        data_info['xml_path'] = img_info['xml_path']

        # deal with xml file
        raw_ann_info = ET.parse(data_info['xml_path'])
        root = raw_ann_info.getroot()
        size = root.find('size')
        if size is not None:
            width = int(size.find('width').text)
            height = int(size.find('height').text)
        else:
            img = Image.open(img_info['img_path'])
            height, width = img.height, img.width
            del img

        data_info['height'] = height
        data_info['width'] = width

        instances = []
        for obj in raw_ann_info.findall('object'):
            instance = {}
            name = obj.find('name').text
            if name not in self._metainfo['classes']:
                continue
            difficult = obj.find('difficult')
            difficult = 0 if difficult is None else int(difficult.text)
            bnd_box = obj.find('bndbox')
            bbox = [
                int(float(bnd_box.find('xmin').text)) - 1,
                int(float(bnd_box.find('ymin').text)) - 1,
                int(float(bnd_box.find('xmax').text)) - 1,
                int(float(bnd_box.find('ymax').text)) - 1
            ]
            instance['ignore_flag'] = difficult
            instance['bbox'] = bbox
            instance['bbox_label'] = self.cat2label[name]
            instances.append(instance)
        data_info['instances'] = instances
        return data_info

    def filter_data(self):
        if self.test_mode:
            return self.data_list

        filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False) \
            if self.filter_cfg is not None else False
        min_size = self.filter_cfg.get('min_size', 0) \
            if self.filter_cfg is not None else 0

        valid_data_infos = []
        for i, data_info in enumerate(self.data_list):
            width = data_info['width']
            height = data_info['height']
            if filter_empty_gt and len(data_info['instances']) == 0:
                continue
            if min(width, height) >= min_size:
                valid_data_infos.append(data_info)

        return valid_data_infos

jittordet/datasets/voc.py
JittorCV/jittordet
from ..engine import BATCH_SAMPLERS, DATASETS
from ..utils import is_list_of
from .base import BaseDetDataset


class PartCompose:

    def __init__(self, composes, partition):
        assert len(composes) == len(partition)
        self.partition = partition
        self.composes = composes

    def __call__(self, data):
        assert 'sample_idx' in data
        sample_idx = data['sample_idx']
        for compose, part in zip(self.composes, self.partition):
            if sample_idx >= part:
                sample_idx -= part
                continue

            return compose(data)


@DATASETS.register_module()
class ConcatDataset(BaseDetDataset):

    def __init__(self,
                 datasets,
                 batch_size,
                 num_workers=0,
                 metainfo=None,
                 test_mode=False,
                 batch_sampler=None,
                 max_refetch=100,
                 **kwargs):
        super(BaseDetDataset, self).__init__(
            batch_size=batch_size, num_workers=num_workers, **kwargs)

        # override some setting in sub dataset
        assert is_list_of(datasets, dict)
        self.data_list, self.lengths = [], []
        transforms = []
        for dataset in datasets:
            dataset['batch_size'] = 1
            dataset['num_workers'] = 0
            dataset['metainfo'] = metainfo
            dataset['test_mode'] = test_mode
            dataset['batch_sampler'] = None
            dataset = DATASETS.build(dataset)
            self.data_list.extend(dataset.data_list)
            self.lengths.append(len(dataset.data_list))
            transforms.append(dataset.transforms)

        self.transforms = PartCompose(transforms, self.lengths)

        self.test_mode = test_mode
        self.max_refetch = max_refetch

        # set total length for jittor.utils.dataset
        self.total_len = len(self.data_list)

        if batch_sampler is not None:
            self.batch_sampler = BATCH_SAMPLERS.build(
                batch_sampler, dataset=self)
        else:
            self.batch_sampler = None

jittordet/datasets/wrappers.py
JittorCV/jittordet
from .config import *  # noqa: F401, F403
from .evaluator import *  # noqa: F401, F403
from .hooks import *  # noqa: F401, F403
from .logger import *  # noqa: F401, F403
from .loops import *  # noqa: F401, F403
from .optim import *  # noqa: F401, F403
from .register import *  # noqa: F401, F403
from .runner import *  # noqa: F401, F403

jittordet/engine/__init__.py
JittorCV/jittordet
from .config import (ConfigDict, ConfigType, MultiConfig, OptConfigType,
                     OptMultiConfig, dump_cfg, load_cfg, merge_cfg)
from .dumpers import cfg_dumpers, json_dumper, yaml_dumper
from .parsers import (cfg_parsers, default_var_parser, env_variable_parser,
                      python_eval_parser, tuple_parser)
from .readers import cfg_readers, yaml_reader
from .utils import delete_node, iter_leaves, set_leaf

__all__ = [
    'load_cfg', 'merge_cfg', 'dump_cfg', 'yaml_reader', 'default_var_parser',
    'cfg_readers', 'env_variable_parser', 'cfg_parsers', 'iter_leaves',
    'set_leaf', 'cfg_dumpers', 'yaml_dumper', 'json_dumper', 'delete_node',
    'ConfigType', 'OptConfigType', 'MultiConfig', 'OptMultiConfig',
    'ConfigDict', 'python_eval_parser', 'tuple_parser'
]

jittordet/engine/config/__init__.py
JittorCV/jittordet
import os
import os.path as osp
from collections.abc import Mapping
from typing import List, Optional, Union

from addict import Dict

from .dumpers import cfg_dumpers
from .parsers import cfg_parsers
from .readers import cfg_readers
from .utils import delete_node

BASE_KEY = '_base_'
COVER_KEY = '_cover_'
RESERVED_KEYS = ['filename']


class ConfigDict(Dict):
    """copy from mmengine https://github.com/open-
    mmlab/mmengine/blob/main/mmengine/config/config.py.

    A dictionary for config which has the same interface as python's built- in
    dictionary and can be used as a normal dictionary. The Config class would
    transform the nested fields (dictionary-like fields) in config file into
    ``ConfigDict``.
    """

    def __missing__(self, name):
        raise KeyError(name)

    def __getattr__(self, name):
        try:
            value = super().__getattr__(name)
        except KeyError:
            raise AttributeError(f"'{self.__class__.__name__}' object has no "
                                 f"attribute '{name}'")
        except Exception as e:
            raise e
        else:
            return value


ConfigType = Union[ConfigDict, dict]
OptConfigType = Optional[ConfigType]

MultiConfig = Union[ConfigType, List[ConfigType]]
OptMultiConfig = Optional[MultiConfig]


def load_cfg(filepath):
    """load cfg from different file."""
    assert osp.isfile(filepath), f'{filepath} is not a exist file'
    ext = osp.splitext(filepath)[-1]
    if ext not in cfg_readers:
        raise NotImplementedError(
            f'Cannot parse "{filepath}" with {ext} type yet')
    cfg = ConfigDict(cfg_readers[ext](filepath))
    for key in RESERVED_KEYS:
        if key in cfg:
            raise KeyError('f"{key}" is a reserved key')

    # use parsers to translate some leaves
    cfg['filename'] = filepath
    for parser in cfg_parsers:
        parser(cfg)

    if BASE_KEY in cfg:
        base_cfg_paths = cfg.pop(BASE_KEY)
        if isinstance(base_cfg_paths, str):
            base_cfg_paths = [base_cfg_paths]
        all_cfg = ConfigDict()
        root_path = osp.dirname(filepath)
        for base_cfg_path in base_cfg_paths:
            if base_cfg_path.startswith('~'):
                base_cfg_path = osp.expanduser(base_cfg_path)
            if base_cfg_path.startswith('.'):
                base_cfg_path = osp.join(root_path, base_cfg_path)
            all_cfg.update(load_cfg(base_cfg_path))
        merge_cfg(all_cfg, cfg)
        cfg = all_cfg

    delete_node(cfg, COVER_KEY)
    return cfg


def merge_cfg(cfg_a, cfg_b):
    """merge cfg_b into cfg_a."""
    for k, v in cfg_b.items():
        if k in cfg_a and (isinstance(v, Mapping)
                           and isinstance(cfg_a[k], Mapping)):
            if v.pop(COVER_KEY, False):
                cfg_a[k] = v
            else:
                merge_cfg(cfg_a[k], v)
        else:
            cfg_a[k] = v


def dump_cfg(cfg, filepath, allow_exist=False, create_dir=True):
    """dump cfg into different files."""
    ext = osp.splitext(filepath)[-1]
    dir_name = osp.dirname(filepath)
    if ext not in cfg_dumpers:
        raise NotImplementedError(f'Cannot dump cfg to {ext} type file yet')

    if osp.exists(filepath) and not allow_exist:
        raise FileExistsError('The target file has existed')

    if dir_name and not osp.exists(dir_name) and create_dir:
        os.makedirs(dir_name)

    cfg_dumpers[ext](cfg, filepath)

jittordet/engine/config/config.py
JittorCV/jittordet
import json

import yaml


def yaml_dumper(cfg, filepath):
    cfg = cfg.to_dict()
    with open(filepath, 'w') as f:
        yaml.dump(cfg, f)


def json_dumper(cfg, filepath):
    cfg = cfg.to_dict()
    with open(filepath, 'w') as f:
        json.dump(cfg, f, indent=4)


cfg_dumpers = {
    '.yml': yaml_dumper,
    '.yaml': yaml_dumper,
    '.json': json_dumper,
}

jittordet/engine/config/dumpers.py
JittorCV/jittordet
import copy
import os
import os.path as osp
import re

from .utils import iter_leaves, set_leaf


def env_variable_parser(cfg):
    """use environment variables in cfg."""
    regexp1 = r'^\s*\$(\w+)\s*\:\s*(\S*?)\s*$'
    regexp2 = r'\<\s*\$(\w+)\s*\:\s*(\S*?)\s*\>'
    for keys, value in iter_leaves(cfg):
        if not isinstance(value, str):
            continue
        # entire string is environment
        results = re.match(regexp1, value)
        if results:
            var_name, def_value = results.groups()
            new_value = os.environ[var_name] \
                if var_name in os.environ else def_value
            if new_value.isdigit():
                new_value = eval(new_value)
            set_leaf(cfg, keys, new_value)
            continue
        # partial string is environment
        results = re.findall(regexp2, value)
        for var_name, def_value in results:
            regexp = r'\<\s*\$' + var_name + r'\s*\:\s*' \
                + def_value + r'\s*\>'
            new_value = os.environ[var_name] if var_name in os.environ \
                else def_value
            value = re.sub(regexp, new_value, value)
        set_leaf(cfg, keys, value)


def default_var_parser(cfg):
    """set some default value in cfg."""
    filename = cfg['filename']
    file_dirname = osp.dirname(filename)
    file_basename = osp.basename(filename)
    file_basename_no_extension = osp.splitext(file_basename)[0]
    file_extname = osp.splitext(filename)[1]
    support_templates = dict(
        fileDirname=file_dirname,
        fileBasename=file_basename,
        fileBasenameNoExtension=file_basename_no_extension,
        fileExtname=file_extname)

    for keys, value in iter_leaves(cfg):
        if not isinstance(value, str):
            continue

        for k, v in support_templates.items():
            regexp = r'\<\s*' + str(k) + r'\s*\>'
            v = v.replace('\\', '/')
            value = re.sub(regexp, v, value)

        set_leaf(cfg, keys, value)


def tuple_parser(cfg):
    for keys, value in iter_leaves(cfg):
        if not isinstance(value, str):
            continue

        if value.startswith('(') and value.endswith(')'):
            try:
                value = eval(value)
            except:  # noqa: E722
                pass

            if isinstance(value, tuple):
                set_leaf(cfg, keys, value)


def python_eval_parser(cfg):
    eval_global = copy.deepcopy(cfg)
    for keys, value in iter_leaves(cfg):
        if not isinstance(value, str):
            continue
        if not value.startswith('<') or not value.endswith('>'):
            continue

        value = value[1:-1]
        value = eval(value, eval_global)
        set_leaf(cfg, keys, value)


cfg_parsers = [
    env_variable_parser, default_var_parser, tuple_parser, python_eval_parser
]

jittordet/engine/config/parsers.py
JittorCV/jittordet
import ast
import json
import os.path as osp
import sys
import types
from importlib import import_module

import yaml


def yaml_reader(filepath):
    with open(filepath, 'r') as f:
        cfg = yaml.load(f, Loader=yaml.Loader)
    return cfg


def json_reader(filepath):
    with open(filepath, 'r') as f:
        cfg = json.load(f)
    return cfg


def python_reader(filepath):
    """Reader python type config.

    Refer to mmcv.utils.config.
    """
    # validate python syntax
    with open(filepath, 'r', encoding='utf-8') as f:
        # Setting encoding explicitly to resolve coding issue on windows
        content = f.read()
    try:
        ast.parse(content)
    except SyntaxError as e:
        raise SyntaxError('There are syntax errors in config '
                          f'file {filepath}: {e}')

    filepath = osp.splitext(filepath)[0]
    dir_name = osp.dirname(filepath)
    module_name = osp.basename(filepath)
    sys.path.insert(0, dir_name)
    mod = import_module(module_name)
    sys.path.pop(0)
    cfg = {
        name: value
        for name, value in mod.__dict__.items() if not name.startswith('__')
        and not isinstance(value, types.ModuleType)
        and not isinstance(value, types.FunctionType)
    }
    # delete imported module
    del sys.modules[module_name]
    return cfg


cfg_readers = {
    '.yml': yaml_reader,
    '.ymal': yaml_reader,
    '.json': json_reader,
    '.py': python_reader
}

jittordet/engine/config/readers.py
JittorCV/jittordet
from collections import Mapping


def iter_leaves(obj):
    """A generator to visit all leaves of obj."""
    if isinstance(obj, Mapping):
        for key, value in obj.items():
            for k, v in iter_leaves(value):
                k.insert(0, key)
                yield (k, v)
    elif isinstance(obj, (list, tuple)):
        for i, value in enumerate(obj):
            for k, v in iter_leaves(value):
                k.insert(0, i)
                yield (k, v)
    else:
        yield [], obj


def set_leaf(obj, keys, value):
    if isinstance(keys, str):
        keys = keys.split('.')
    for k in keys[:-1]:
        obj = obj[k]
    obj[keys[-1]] = value


def delete_node(obj, keys):
    if isinstance(keys, (tuple, list)):
        for key in keys[:-1]:
            obj = obj[key]
        del obj[keys[-1]]
    else:
        assert isinstance(keys, str), 'only support search str node'
        if isinstance(obj, Mapping):
            if keys in obj:
                del obj[keys]
            for value in obj.values():
                delete_node(value, keys)
        elif isinstance(obj, (tuple, list)):
            for value in obj:
                delete_node(value, keys)

jittordet/engine/config/utils.py
JittorCV/jittordet
from .base_evaluator import BaseEvaluator
from .coco_evaluator import CocoEvaluator
from .voc_evaluator import VocEvaluator

__all__ = ['BaseEvaluator', 'CocoEvaluator', 'VocEvaluator']

jittordet/engine/evaluator/__init__.py
JittorCV/jittordet
import os
import os.path as osp
import pickle
import shutil
import tempfile
from abc import ABCMeta, abstractmethod

import jittor as jt


class BaseEvaluator(metaclass=ABCMeta):

    @abstractmethod
    def process(self, dataset, data_samples):
        pass

    @abstractmethod
    def compute_metrics(self, dataset, results, logger):
        pass

    def evaluate(self, dataset, logger):
        results = self.collect_results(self.results, len(dataset.data_list))
        if not jt.in_mpi or jt.rank == 0:
            metrics = self.compute_metrics(dataset, results, logger)
        else:
            metrics = None
        metrics = self.broadcast_metrics(metrics)
        self.results.clear()
        return metrics

    def gen_broadcasted_tmpdir(self):
        MAX_LEN = 512
        # 32 is whitespace
        dir_tensor = jt.full((MAX_LEN, ), 32, dtype=jt.uint8)
        if jt.rank == 0:
            if not osp.exists('.dist_test'):
                os.makedirs('.dist_test')
            tmpdir = tempfile.mkdtemp(dir='.dist_test')
            tmpdir = jt.array(bytearray(tmpdir.encode()), dtype=jt.uint8)
            dir_tensor[:len(tmpdir)] = tmpdir
        dir_tensor = dir_tensor.mpi_broadcast(root=0)
        tmpdir = dir_tensor.numpy().tobytes().decode().rstrip()
        return tmpdir

    def collect_results(self, result_part, size):
        """Collect results under cpu mode."""
        rank, world_size = jt.rank, jt.world_size
        if world_size == 1:
            return result_part[:size]

        tmpdir = self.gen_broadcasted_tmpdir()
        # dump the part result to the dir
        with open(osp.join(tmpdir, f'part_{rank}.pkl'),
                  'wb') as f:  # type: ignore
            pickle.dump(result_part, f, protocol=2)

        self.barrier()

        if rank != 0:
            return None
        else:
            # load results of all parts from tmp dir
            part_list = []
            for i in range(world_size):
                path = osp.join(tmpdir, f'part_{i}.pkl')  # type: ignore
                with open(path, 'rb') as f:
                    part_list.append(pickle.load(f))
            # sort the results
            ordered_results = []
            for res in zip(*part_list):
                ordered_results.extend(list(res))
            # the dataloader may pad some samples
            ordered_results = ordered_results[:size]
            # remove tmp dir
            shutil.rmtree(tmpdir)  # type: ignore
            return ordered_results

    def broadcast_metrics(self, metrics):
        if jt.world_size == 1:
            return metrics

        tmpdir = self.gen_broadcasted_tmpdir()
        if jt.rank == 0:
            with open(osp.join(tmpdir, 'metrics.pkl'), 'wb') as f:
                pickle.dump(metrics, f, protocol=2)

        self.barrier()

        with open(osp.join(tmpdir, 'metrics.pkl'), 'rb') as f:
            metrics = pickle.load(f)

        if jt.rank == 0:
            shutil.rmtree(tmpdir)  # type: ignore
        return metrics

    @staticmethod
    def barrier():
        if jt.in_mpi:
            lock = jt.array([1])
            lock = lock.mpi_all_reduce('mean')
            lock.sync(device_sync=True)
            del lock

jittordet/engine/evaluator/base_evaluator.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import datetime
import itertools
import json
import os.path as osp
import tempfile
from collections import OrderedDict
from typing import Dict, List, Optional, Sequence, Union

import numpy as np
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from terminaltables import AsciiTable

from jittordet.structures import BaseDataElement
from ..register import EVALUATORS
from .base_evaluator import BaseEvaluator
from .utils import eval_recalls


@EVALUATORS.register_module()
class CocoEvaluator(BaseEvaluator):
    """COCO evaluation metric.

    Evaluate AR, AP, and mAP for detection tasks including proposal/box
    detection and instance segmentation. Please refer to
    https://cocodataset.org/#detection-eval for more details.

    Args:
        ann_file (str, optional): Path to the coco format annotation file.
            If not specified, ground truth annotations from the dataset will
            be converted to coco format. Defaults to None.
        metric (str | List[str]): Metrics to be evaluated. Valid metrics
            include 'bbox', 'segm', 'proposal', and 'proposal_fast'.
            Defaults to 'bbox'.
        classwise (bool): Whether to evaluate the metric class-wise.
            Defaults to False.
        proposal_nums (Sequence[int]): Numbers of proposals to be evaluated.
            Defaults to (100, 300, 1000).
        iou_thrs (float | List[float], optional): IoU threshold to compute AP
            and AR. If not specified, IoUs from 0.5 to 0.95 will be used.
            Defaults to None.
        metric_items (List[str], optional): Metric result names to be
            recorded in the evaluation result. Defaults to None.
        format_only (bool): Format the output results without perform
            evaluation. It is useful when you want to format the result
            to a specific format and submit it to the test server.
            Defaults to False.
        outfile_prefix (str, optional): The prefix of json files. It includes
            the file path and the prefix of filename, e.g., "a/b/prefix".
            If not specified, a temp file will be created. Defaults to None.
        file_client_args (dict): Arguments to instantiate a FileClient.
            See :class:`mmengine.fileio.FileClient` for details.
            Defaults to ``dict(backend='disk')``.
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
            'gpu'. Defaults to 'cpu'.
        prefix (str, optional): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
            If prefix is not provided in the argument, self.default_prefix
            will be used instead. Defaults to None.
    """

    def __init__(self,
                 ann_file: Optional[str] = None,
                 metric: Union[str, List[str]] = 'bbox',
                 classwise: bool = False,
                 proposal_nums: Sequence[int] = (100, 300, 1000),
                 iou_thrs: Optional[Union[float, Sequence[float]]] = None,
                 metric_items: Optional[Sequence[str]] = None,
                 format_only: bool = False,
                 outfile_prefix: Optional[str] = None) -> None:
        self.results = []
        # coco evaluation metrics
        self.metrics = metric if isinstance(metric, list) else [metric]
        allowed_metrics = ['bbox', 'segm', 'proposal', 'proposal_fast']
        for metric in self.metrics:
            if metric not in allowed_metrics:
                raise KeyError(
                    "metric should be one of 'bbox', 'segm', 'proposal', "
                    f"'proposal_fast', but got {metric}.")

        # do class wise evaluation, default False
        self.classwise = classwise

        # proposal_nums used to compute recall or precision.
        self.proposal_nums = list(proposal_nums)

        # iou_thrs used to compute recall or precision.
        if iou_thrs is None:
            iou_thrs = np.linspace(
                .5, 0.95, int(np.round((0.95 - .5) / .05)) + 1, endpoint=True)
        self.iou_thrs = iou_thrs
        self.metric_items = metric_items
        self.format_only = format_only
        if self.format_only:
            assert outfile_prefix is not None, 'outfile_prefix must be not'
            'None when format_only is True, otherwise the result files will'
            'be saved to a temp directory which will be cleaned up at the end.'

        self.outfile_prefix = outfile_prefix

        # if ann_file is not specified,
        # initialize coco api with the converted dataset
        if ann_file is not None:
            self._coco_api = COCO(ann_file)
        else:
            self._coco_api = None

        # handle dataset lazy init
        self.cat_ids = None
        self.img_ids = None

    def fast_eval_recall(self,
                         results: List[dict],
                         proposal_nums: Sequence[int],
                         iou_thrs: Sequence[float],
                         logger=None) -> np.ndarray:
        """Evaluate proposal recall with COCO's fast_eval_recall.

        Args:
            results (List[dict]): Results of the dataset.
            proposal_nums (Sequence[int]): Proposal numbers used for
                evaluation.
            iou_thrs (Sequence[float]): IoU thresholds used for evaluation.
            logger (MMLogger, optional): Logger used for logging the recall
                summary.
        Returns:
            np.ndarray: Averaged recall results.
        """
        gt_bboxes = []
        pred_bboxes = [result['bboxes'] for result in results]
        for i in range(len(self.img_ids)):
            ann_ids = self._coco_api.getAnnIds(imgIds=self.img_ids[i])
            ann_info = self._coco_api.loadAnns(ann_ids)
            if len(ann_info) == 0:
                gt_bboxes.append(np.zeros((0, 4)))
                continue
            bboxes = []
            for ann in ann_info:
                if ann.get('ignore', False) or ann['iscrowd']:
                    continue
                x1, y1, w, h = ann['bbox']
                bboxes.append([x1, y1, x1 + w, y1 + h])
            bboxes = np.array(bboxes, dtype=np.float32)
            if bboxes.shape[0] == 0:
                bboxes = np.zeros((0, 4))
            gt_bboxes.append(bboxes)

        recalls = eval_recalls(
            gt_bboxes, pred_bboxes, proposal_nums, iou_thrs, logger=logger)
        ar = recalls.mean(axis=1)
        return ar

    def xyxy2xywh(self, bbox: np.ndarray) -> list:
        """Convert ``xyxy`` style bounding boxes to ``xywh`` style for COCO
        evaluation.

        Args:
            bbox (numpy.ndarray): The bounding boxes, shape (4, ), in
                ``xyxy`` order.

        Returns:
            list[float]: The converted bounding boxes, in ``xywh`` order.
        """

        _bbox: List = bbox.tolist()
        return [
            _bbox[0],
            _bbox[1],
            _bbox[2] - _bbox[0],
            _bbox[3] - _bbox[1],
        ]

    def results2json(self, results: Sequence[dict],
                     outfile_prefix: str) -> dict:
        """Dump the detection results to a COCO style json file.

        There are 3 types of results: proposals, bbox predictions, mask
        predictions, and they have different data types. This method will
        automatically recognize the type, and dump them to json files.

        Args:
            results (Sequence[dict]): Testing results of the
                dataset.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json files will be named
                "somepath/xxx.bbox.json", "somepath/xxx.segm.json",
                "somepath/xxx.proposal.json".

        Returns:
            dict: Possible keys are "bbox", "segm", "proposal", and
            values are corresponding filenames.
        """
        bbox_json_results = []
        segm_json_results = [] if 'masks' in results[0] else None
        for idx, result in enumerate(results):
            image_id = result.get('img_id', idx)
            labels = result['labels']
            bboxes = result['bboxes']
            scores = result['scores']
            # bbox results
            for i, label in enumerate(labels):
                data = dict()
                data['image_id'] = image_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(scores[i])
                data['category_id'] = self.cat_ids[label]
                bbox_json_results.append(data)

            if segm_json_results is None:
                continue

            # segm results
            masks = result['masks']
            mask_scores = result.get('mask_scores', scores)
            for i, label in enumerate(labels):
                data = dict()
                data['image_id'] = image_id
                data['bbox'] = self.xyxy2xywh(bboxes[i])
                data['score'] = float(mask_scores[i])
                data['category_id'] = self.cat_ids[label]
                if isinstance(masks[i]['counts'], bytes):
                    masks[i]['counts'] = masks[i]['counts'].decode()
                data['segmentation'] = masks[i]
                segm_json_results.append(data)

        result_files = dict()
        result_files['bbox'] = f'{outfile_prefix}.bbox.json'
        result_files['proposal'] = f'{outfile_prefix}.bbox.json'
        with open(result_files['bbox'], 'w') as f:
            json.dump(bbox_json_results, f)

        if segm_json_results is not None:
            result_files['segm'] = f'{outfile_prefix}.segm.json'
            with open(result_files['segm'], 'w') as f:
                json.dump(segm_json_results, f)

        return result_files

    def gt_to_coco_json(self, gt_dicts: Sequence[dict], outfile_prefix: str,
                        class_names: List[str]) -> str:
        """Convert ground truth to coco format json file.

        Args:
            gt_dicts (Sequence[dict]): Ground truth of the dataset.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json file will be named
                "somepath/xxx.gt.json".
        Returns:
            str: The filename of the json file.
        """
        categories = [
            dict(id=id, name=name) for id, name in enumerate(class_names)
        ]
        image_infos = []
        annotations = []

        for idx, gt_dict in enumerate(gt_dicts):
            img_id = gt_dict.get('img_id', idx)
            image_info = dict(
                id=img_id,
                width=gt_dict['width'],
                height=gt_dict['height'],
                file_name='')
            image_infos.append(image_info)
            for ann in gt_dict['instances']:
                label = ann['bbox_label']
                bbox = ann['bbox']
                coco_bbox = [
                    bbox[0],
                    bbox[1],
                    bbox[2] - bbox[0],
                    bbox[3] - bbox[1],
                ]

                annotation = dict(
                    id=len(annotations) +
                    1,  # coco api requires id starts with 1
                    image_id=img_id,
                    bbox=coco_bbox,
                    iscrowd=ann.get('ignore_flag', 0),
                    category_id=int(label),
                    area=coco_bbox[2] * coco_bbox[3])
                if ann.get('mask', None):
                    mask = ann['mask']
                    # area = mask_util.area(mask)
                    if isinstance(mask, dict) and isinstance(
                            mask['counts'], bytes):
                        mask['counts'] = mask['counts'].decode()
                    annotation['segmentation'] = mask
                    # annotation['area'] = float(area)
                annotations.append(annotation)

        info = dict(
            date_created=str(datetime.datetime.now()),
            description='Coco json file converted by mmdet CocoMetric.')
        coco_json = dict(
            info=info,
            images=image_infos,
            categories=categories,
            licenses=None,
        )
        if len(annotations) > 0:
            coco_json['annotations'] = annotations
        converted_json_path = f'{outfile_prefix}.gt.json'
        with open(converted_json_path, 'w') as f:
            json.dump(coco_json, f)
        return converted_json_path

    def process(self, dataset, data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and predictions. The processed
        results should be stored in ``self.results``, which will be used to
        compute the metrics when all batches have been processed.

        Args:
            data_batch (dict): A batch of data from the dataloader.
            data_samples (Sequence[dict]): A batch of data samples that
                contain annotations and predictions.
        """
        for data_sample in data_samples:
            if isinstance(data_sample, BaseDataElement):
                data_sample = data_sample.to_dict()

            result = dict()
            pred = data_sample['pred_instances']
            result['img_id'] = data_sample['img_id']
            result['bboxes'] = pred['bboxes'].numpy()
            result['scores'] = pred['scores'].numpy()
            result['labels'] = pred['labels'].numpy()
            # TODO: enable instance segmentation function.
            # encode mask to RLE
            # if 'masks' in pred:
            #     result['masks'] = encode_mask_results(
            #         pred['masks'].detach().cpu().numpy())
            # some detectors use different scores for bbox and mask
            if 'mask_scores' in pred:
                result['mask_scores'] = pred['mask_scores'].cpu().numpy()

            # parse gt
            if self._coco_api is None:
                sample_idx = data_sample['sample_idx']
                gt = copy.deepcopy(dataset.data_list[sample_idx])
            else:
                gt = None
            # add converted result to the results list
            self.results.append((gt, result))

    def compute_metrics(self, dataset, results: list,
                        logger) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        # split gt and prediction list
        gts, preds = zip(*results)

        tmp_dir = None
        if self.outfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            outfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            outfile_prefix = self.outfile_prefix

        if self._coco_api is None:
            # use converted gt json file to initialize coco api
            logger.info('Converting ground truth to coco format...')
            coco_json_path = self.gt_to_coco_json(
                gt_dicts=gts,
                outfile_prefix=outfile_prefix,
                class_names=dataset.metainfo['classes'])
            self._coco_api = COCO(coco_json_path)

        # handle lazy init
        if self.cat_ids is None:
            self.cat_ids = self._coco_api.getCatIds(
                catNms=dataset.metainfo['classes'])
        if self.img_ids is None:
            self.img_ids = self._coco_api.getImgIds()

        # convert predictions to coco format and dump to json file
        result_files = self.results2json(preds, outfile_prefix)

        eval_results = OrderedDict()
        if self.format_only:
            logger.info('results are saved in '
                        f'{osp.dirname(outfile_prefix)}')
            return eval_results

        for metric in self.metrics:
            logger.info(f'Evaluating {metric}...')

            # TODO: May refactor fast_eval_recall to an independent metric?
            # fast eval recall
            if metric == 'proposal_fast':
                ar = self.fast_eval_recall(
                    preds, self.proposal_nums, self.iou_thrs, logger=logger)
                log_msg = []
                for i, num in enumerate(self.proposal_nums):
                    eval_results[f'AR@{num}'] = ar[i]
                    log_msg.append(f'\nAR@{num}\t{ar[i]:.4f}')
                log_msg = ''.join(log_msg)
                logger.info(log_msg)
                continue

            # evaluate proposal, bbox and segm
            iou_type = 'bbox' if metric == 'proposal' else metric
            if metric not in result_files:
                raise KeyError(f'{metric} is not in results')
            try:
                with open(result_files[metric], 'r') as f:
                    predictions = json.load(f)
                if iou_type == 'segm':
                    # Refer to https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/coco.py#L331  # noqa
                    # When evaluating mask AP, if the results contain bbox,
                    # cocoapi will use the box area instead of the mask area
                    # for calculating the instance area. Though the overall AP
                    # is not affected, this leads to different
                    # small/medium/large mask AP results.
                    for x in predictions:
                        x.pop('bbox')
                coco_dt = self._coco_api.loadRes(predictions)

            except IndexError:
                logger.error(
                    'The testing results of the whole dataset is empty.')
                break

            coco_eval = COCOeval(self._coco_api, coco_dt, iou_type)

            coco_eval.params.catIds = self.cat_ids
            coco_eval.params.imgIds = self.img_ids
            coco_eval.params.maxDets = list(self.proposal_nums)
            coco_eval.params.iouThrs = self.iou_thrs

            # mapping of cocoEval.stats
            coco_metric_names = {
                'mAP': 0,
                'mAP_50': 1,
                'mAP_75': 2,
                'mAP_s': 3,
                'mAP_m': 4,
                'mAP_l': 5,
                'AR@100': 6,
                'AR@300': 7,
                'AR@1000': 8,
                'AR_s@1000': 9,
                'AR_m@1000': 10,
                'AR_l@1000': 11
            }
            metric_items = self.metric_items
            if metric_items is not None:
                for metric_item in metric_items:
                    if metric_item not in coco_metric_names:
                        raise KeyError(
                            f'metric item "{metric_item}" is not supported')

            if metric == 'proposal':
                coco_eval.params.useCats = 0
                coco_eval.evaluate()
                coco_eval.accumulate()
                coco_eval.summarize()
                if metric_items is None:
                    metric_items = [
                        'AR@100', 'AR@300', 'AR@1000', 'AR_s@1000',
                        'AR_m@1000', 'AR_l@1000'
                    ]

                for item in metric_items:
                    val = float(
                        f'{coco_eval.stats[coco_metric_names[item]]:.3f}')
                    eval_results[item] = val
            else:
                coco_eval.evaluate()
                coco_eval.accumulate()
                coco_eval.summarize()
                if self.classwise:  # Compute per-category AP
                    # Compute per-category AP
                    # from https://github.com/facebookresearch/detectron2/
                    precisions = coco_eval.eval['precision']
                    # precision: (iou, recall, cls, area range, max dets)
                    assert len(self.cat_ids) == precisions.shape[2]

                    results_per_category = []
                    for idx, cat_id in enumerate(self.cat_ids):
                        # area range index 0: all area ranges
                        # max dets index -1: typically 100 per image
                        nm = self._coco_api.loadCats(cat_id)[0]
                        precision = precisions[:, :, idx, 0, -1]
                        precision = precision[precision > -1]
                        if precision.size:
                            ap = np.mean(precision)
                        else:
                            ap = float('nan')
                        results_per_category.append(
                            (f'{nm["name"]}', f'{round(ap, 3)}'))
                        eval_results[f'{nm["name"]}_precision'] = round(ap, 3)

                    num_columns = min(6, len(results_per_category) * 2)
                    results_flatten = list(
                        itertools.chain(*results_per_category))
                    headers = ['category', 'AP'] * (num_columns // 2)
                    results_2d = itertools.zip_longest(*[
                        results_flatten[i::num_columns]
                        for i in range(num_columns)
                    ])
                    table_data = [headers]
                    table_data += [result for result in results_2d]
                    table = AsciiTable(table_data)
                    logger.info('\n' + table.table)

                if metric_items is None:
                    metric_items = [
                        'mAP', 'mAP_50', 'mAP_75', 'mAP_s', 'mAP_m', 'mAP_l'
                    ]

                for metric_item in metric_items:
                    key = f'{metric}_{metric_item}'
                    val = coco_eval.stats[coco_metric_names[metric_item]]
                    eval_results[key] = float(f'{round(val, 3)}')

                ap = coco_eval.stats[:6]
                # logger.info(f'{metric}_mAP_copypaste: {ap[0]:.3f} '
                #             f'{ap[1]:.3f} {ap[2]:.3f} {ap[3]:.3f} '
                #             f'{ap[4]:.3f} {ap[5]:.3f}')

        if tmp_dir is not None:
            tmp_dir.cleanup()
        return eval_results

jittordet/engine/evaluator/coco_evaluator.py
JittorCV/jittordet
from collections.abc import Sequence
from multiprocessing import Pool

import numpy as np
from terminaltables import AsciiTable

from ..logger import print_log


def bbox_overlaps(bboxes1,
                  bboxes2,
                  mode='iou',
                  eps=1e-6,
                  use_legacy_coordinate=False):
    """Calculate the ious between each bbox of bboxes1 and bboxes2."""

    assert mode in ['iou', 'iof']
    if not use_legacy_coordinate:
        extra_length = 0.
    else:
        extra_length = 1.
    bboxes1 = bboxes1.astype(np.float32)
    bboxes2 = bboxes2.astype(np.float32)
    rows = bboxes1.shape[0]
    cols = bboxes2.shape[0]
    ious = np.zeros((rows, cols), dtype=np.float32)
    if rows * cols == 0:
        return ious
    exchange = False
    if bboxes1.shape[0] > bboxes2.shape[0]:
        bboxes1, bboxes2 = bboxes2, bboxes1
        ious = np.zeros((cols, rows), dtype=np.float32)
        exchange = True
    area1 = (bboxes1[:, 2] - bboxes1[:, 0] + extra_length) * (
        bboxes1[:, 3] - bboxes1[:, 1] + extra_length)
    area2 = (bboxes2[:, 2] - bboxes2[:, 0] + extra_length) * (
        bboxes2[:, 3] - bboxes2[:, 1] + extra_length)
    for i in range(bboxes1.shape[0]):
        x_start = np.maximum(bboxes1[i, 0], bboxes2[:, 0])
        y_start = np.maximum(bboxes1[i, 1], bboxes2[:, 1])
        x_end = np.minimum(bboxes1[i, 2], bboxes2[:, 2])
        y_end = np.minimum(bboxes1[i, 3], bboxes2[:, 3])
        overlap = np.maximum(x_end - x_start + extra_length, 0) * np.maximum(
            y_end - y_start + extra_length, 0)
        if mode == 'iou':
            union = area1[i] + area2 - overlap
        else:
            union = area1[i] if not exchange else area2
        union = np.maximum(union, eps)
        ious[i, :] = overlap / union
    if exchange:
        ious = ious.T
    return ious


def average_precision(recalls, precisions, mode='area'):
    """Calculate average precision (for single or multiple scales)."""
    no_scale = False
    if recalls.ndim == 1:
        no_scale = True
        recalls = recalls[np.newaxis, :]
        precisions = precisions[np.newaxis, :]
    assert recalls.shape == precisions.shape and recalls.ndim == 2
    num_scales = recalls.shape[0]
    ap = np.zeros(num_scales, dtype=np.float32)
    if mode == 'area':
        zeros = np.zeros((num_scales, 1), dtype=recalls.dtype)
        ones = np.ones((num_scales, 1), dtype=recalls.dtype)
        mrec = np.hstack((zeros, recalls, ones))
        mpre = np.hstack((zeros, precisions, zeros))
        for i in range(mpre.shape[1] - 1, 0, -1):
            mpre[:, i - 1] = np.maximum(mpre[:, i - 1], mpre[:, i])
        for i in range(num_scales):
            ind = np.where(mrec[i, 1:] != mrec[i, :-1])[0]
            ap[i] = np.sum(
                (mrec[i, ind + 1] - mrec[i, ind]) * mpre[i, ind + 1])
    elif mode == '11points':
        for i in range(num_scales):
            for thr in np.arange(0, 1 + 1e-3, 0.1):
                precs = precisions[i, recalls[i, :] >= thr]
                prec = precs.max() if precs.size > 0 else 0
                ap[i] += prec
        ap /= 11
    else:
        raise ValueError(
            'Unrecognized mode, only "area" and "11points" are supported')
    if no_scale:
        ap = ap[0]
    return ap


def tpfp_imagenet(det_bboxes,
                  gt_bboxes,
                  gt_bboxes_ignore=None,
                  default_iou_thr=0.5,
                  area_ranges=None,
                  use_legacy_coordinate=False,
                  **kwargs):
    """Check if detected bboxes are true positive or false positive."""
    if not use_legacy_coordinate:
        extra_length = 0.
    else:
        extra_length = 1.

    # an indicator of ignored gts
    gt_ignore_inds = np.concatenate(
        (np.zeros(gt_bboxes.shape[0], dtype=np.bool),
         np.ones(gt_bboxes_ignore.shape[0], dtype=np.bool)))
    # stack gt_bboxes and gt_bboxes_ignore for convenience
    gt_bboxes = np.vstack((gt_bboxes, gt_bboxes_ignore))

    num_dets = det_bboxes.shape[0]
    num_gts = gt_bboxes.shape[0]
    if area_ranges is None:
        area_ranges = [(None, None)]
    num_scales = len(area_ranges)
    # tp and fp are of shape (num_scales, num_gts), each row is tp or fp
    # of a certain scale.
    tp = np.zeros((num_scales, num_dets), dtype=np.float32)
    fp = np.zeros((num_scales, num_dets), dtype=np.float32)
    if gt_bboxes.shape[0] == 0:
        if area_ranges == [(None, None)]:
            fp[...] = 1
        else:
            det_areas = (
                det_bboxes[:, 2] - det_bboxes[:, 0] + extra_length) * (
                    det_bboxes[:, 3] - det_bboxes[:, 1] + extra_length)
            for i, (min_area, max_area) in enumerate(area_ranges):
                fp[i, (det_areas >= min_area) & (det_areas < max_area)] = 1
        return tp, fp
    ious = bbox_overlaps(
        det_bboxes, gt_bboxes - 1, use_legacy_coordinate=use_legacy_coordinate)
    gt_w = gt_bboxes[:, 2] - gt_bboxes[:, 0] + extra_length
    gt_h = gt_bboxes[:, 3] - gt_bboxes[:, 1] + extra_length
    iou_thrs = np.minimum((gt_w * gt_h) / ((gt_w + 10.0) * (gt_h + 10.0)),
                          default_iou_thr)
    # sort all detections by scores in descending order
    sort_inds = np.argsort(-det_bboxes[:, -1])
    for k, (min_area, max_area) in enumerate(area_ranges):
        gt_covered = np.zeros(num_gts, dtype=bool)
        # if no area range is specified, gt_area_ignore is all False
        if min_area is None:
            gt_area_ignore = np.zeros_like(gt_ignore_inds, dtype=bool)
        else:
            gt_areas = gt_w * gt_h
            gt_area_ignore = (gt_areas < min_area) | (gt_areas >= max_area)
        for i in sort_inds:
            max_iou = -1
            matched_gt = -1
            # find best overlapped available gt
            for j in range(num_gts):
                # different from PASCAL VOC: allow finding other gts if the
                # best overlapped ones are already matched by other det bboxes
                if gt_covered[j]:
                    continue
                elif ious[i, j] >= iou_thrs[j] and ious[i, j] > max_iou:
                    max_iou = ious[i, j]
                    matched_gt = j
            # there are 4 cases for a det bbox:
            # 1. it matches a gt, tp = 1, fp = 0
            # 2. it matches an ignored gt, tp = 0, fp = 0
            # 3. it matches no gt and within area range, tp = 0, fp = 1
            # 4. it matches no gt but is beyond area range, tp = 0, fp = 0
            if matched_gt >= 0:
                gt_covered[matched_gt] = 1
                if not (gt_ignore_inds[matched_gt]
                        or gt_area_ignore[matched_gt]):
                    tp[k, i] = 1
            elif min_area is None:
                fp[k, i] = 1
            else:
                bbox = det_bboxes[i, :4]
                area = (bbox[2] - bbox[0] + extra_length) * (
                    bbox[3] - bbox[1] + extra_length)
                if area >= min_area and area < max_area:
                    fp[k, i] = 1
    return tp, fp


def tpfp_default(det_bboxes,
                 gt_bboxes,
                 gt_bboxes_ignore=None,
                 iou_thr=0.5,
                 area_ranges=None,
                 use_legacy_coordinate=False,
                 **kwargs):
    """Check if detected bboxes are true positive or false positive."""

    if not use_legacy_coordinate:
        extra_length = 0.
    else:
        extra_length = 1.

    # an indicator of ignored gts
    gt_ignore_inds = np.concatenate(
        (np.zeros(gt_bboxes.shape[0], dtype=np.bool),
         np.ones(gt_bboxes_ignore.shape[0], dtype=np.bool)))
    # stack gt_bboxes and gt_bboxes_ignore for convenience
    gt_bboxes = np.vstack((gt_bboxes, gt_bboxes_ignore))

    num_dets = det_bboxes.shape[0]
    num_gts = gt_bboxes.shape[0]
    if area_ranges is None:
        area_ranges = [(None, None)]
    num_scales = len(area_ranges)
    # tp and fp are of shape (num_scales, num_gts), each row is tp or fp of
    # a certain scale
    tp = np.zeros((num_scales, num_dets), dtype=np.float32)
    fp = np.zeros((num_scales, num_dets), dtype=np.float32)

    # if there is no gt bboxes in this image, then all det bboxes
    # within area range are false positives
    if gt_bboxes.shape[0] == 0:
        if area_ranges == [(None, None)]:
            fp[...] = 1
        else:
            det_areas = (
                det_bboxes[:, 2] - det_bboxes[:, 0] + extra_length) * (
                    det_bboxes[:, 3] - det_bboxes[:, 1] + extra_length)
            for i, (min_area, max_area) in enumerate(area_ranges):
                fp[i, (det_areas >= min_area) & (det_areas < max_area)] = 1
        return tp, fp

    ious = bbox_overlaps(
        det_bboxes, gt_bboxes, use_legacy_coordinate=use_legacy_coordinate)
    # for each det, the max iou with all gts
    ious_max = ious.max(axis=1)
    # for each det, which gt overlaps most with it
    ious_argmax = ious.argmax(axis=1)
    # sort all dets in descending order by scores
    sort_inds = np.argsort(-det_bboxes[:, -1])
    for k, (min_area, max_area) in enumerate(area_ranges):
        gt_covered = np.zeros(num_gts, dtype=bool)
        # if no area range is specified, gt_area_ignore is all False
        if min_area is None:
            gt_area_ignore = np.zeros_like(gt_ignore_inds, dtype=bool)
        else:
            gt_areas = (gt_bboxes[:, 2] - gt_bboxes[:, 0] + extra_length) * (
                gt_bboxes[:, 3] - gt_bboxes[:, 1] + extra_length)
            gt_area_ignore = (gt_areas < min_area) | (gt_areas >= max_area)
        for i in sort_inds:
            if ious_max[i] >= iou_thr:
                matched_gt = ious_argmax[i]
                if not (gt_ignore_inds[matched_gt]
                        or gt_area_ignore[matched_gt]):
                    if not gt_covered[matched_gt]:
                        gt_covered[matched_gt] = True
                        tp[k, i] = 1
                    else:
                        fp[k, i] = 1
                # otherwise ignore this detected bbox, tp = 0, fp = 0
            elif min_area is None:
                fp[k, i] = 1
            else:
                bbox = det_bboxes[i, :4]
                area = (bbox[2] - bbox[0] + extra_length) * (
                    bbox[3] - bbox[1] + extra_length)
                if area >= min_area and area < max_area:
                    fp[k, i] = 1
    return tp, fp


def tpfp_openimages(det_bboxes,
                    gt_bboxes,
                    gt_bboxes_ignore=None,
                    iou_thr=0.5,
                    area_ranges=None,
                    use_legacy_coordinate=False,
                    gt_bboxes_group_of=None,
                    use_group_of=True,
                    ioa_thr=0.5,
                    **kwargs):
    """Check if detected bboxes are true positive or false positive."""

    if not use_legacy_coordinate:
        extra_length = 0.
    else:
        extra_length = 1.

    # an indicator of ignored gts
    gt_ignore_inds = np.concatenate(
        (np.zeros(gt_bboxes.shape[0], dtype=np.bool),
         np.ones(gt_bboxes_ignore.shape[0], dtype=np.bool)))
    # stack gt_bboxes and gt_bboxes_ignore for convenience
    gt_bboxes = np.vstack((gt_bboxes, gt_bboxes_ignore))

    num_dets = det_bboxes.shape[0]
    num_gts = gt_bboxes.shape[0]
    if area_ranges is None:
        area_ranges = [(None, None)]
    num_scales = len(area_ranges)
    # tp and fp are of shape (num_scales, num_gts), each row is tp or fp of
    # a certain scale
    tp = np.zeros((num_scales, num_dets), dtype=np.float32)
    fp = np.zeros((num_scales, num_dets), dtype=np.float32)

    # if there is no gt bboxes in this image, then all det bboxes
    # within area range are false positives
    if gt_bboxes.shape[0] == 0:
        if area_ranges == [(None, None)]:
            fp[...] = 1
        else:
            det_areas = (
                det_bboxes[:, 2] - det_bboxes[:, 0] + extra_length) * (
                    det_bboxes[:, 3] - det_bboxes[:, 1] + extra_length)
            for i, (min_area, max_area) in enumerate(area_ranges):
                fp[i, (det_areas >= min_area) & (det_areas < max_area)] = 1
        return tp, fp, det_bboxes

    if gt_bboxes_group_of is not None and use_group_of:
        # if handle group-of boxes, divided gt boxes into two parts:
        # non-group-of and group-of.Then calculate ious and ioas through
        # non-group-of group-of gts respectively. This only used in
        # OpenImages evaluation.
        assert gt_bboxes_group_of.shape[0] == gt_bboxes.shape[0]
        non_group_gt_bboxes = gt_bboxes[~gt_bboxes_group_of]
        group_gt_bboxes = gt_bboxes[gt_bboxes_group_of]
        num_gts_group = group_gt_bboxes.shape[0]
        ious = bbox_overlaps(det_bboxes, non_group_gt_bboxes)
        ioas = bbox_overlaps(det_bboxes, group_gt_bboxes, mode='iof')
    else:
        # if not consider group-of boxes, only calculate ious through gt boxes
        ious = bbox_overlaps(
            det_bboxes, gt_bboxes, use_legacy_coordinate=use_legacy_coordinate)
        ioas = None

    if ious.shape[1] > 0:
        # for each det, the max iou with all gts
        ious_max = ious.max(axis=1)
        # for each det, which gt overlaps most with it
        ious_argmax = ious.argmax(axis=1)
        # sort all dets in descending order by scores
        sort_inds = np.argsort(-det_bboxes[:, -1])
        for k, (min_area, max_area) in enumerate(area_ranges):
            gt_covered = np.zeros(num_gts, dtype=bool)
            # if no area range is specified, gt_area_ignore is all False
            if min_area is None:
                gt_area_ignore = np.zeros_like(gt_ignore_inds, dtype=bool)
            else:
                gt_areas = (
                    gt_bboxes[:, 2] - gt_bboxes[:, 0] + extra_length) * (
                        gt_bboxes[:, 3] - gt_bboxes[:, 1] + extra_length)
                gt_area_ignore = (gt_areas < min_area) | (gt_areas >= max_area)
            for i in sort_inds:
                if ious_max[i] >= iou_thr:
                    matched_gt = ious_argmax[i]
                    if not (gt_ignore_inds[matched_gt]
                            or gt_area_ignore[matched_gt]):
                        if not gt_covered[matched_gt]:
                            gt_covered[matched_gt] = True
                            tp[k, i] = 1
                        else:
                            fp[k, i] = 1
                    # otherwise ignore this detected bbox, tp = 0, fp = 0
                elif min_area is None:
                    fp[k, i] = 1
                else:
                    bbox = det_bboxes[i, :4]
                    area = (bbox[2] - bbox[0] + extra_length) * (
                        bbox[3] - bbox[1] + extra_length)
                    if area >= min_area and area < max_area:
                        fp[k, i] = 1
    else:
        # if there is no no-group-of gt bboxes in this image,
        # then all det bboxes within area range are false positives.
        # Only used in OpenImages evaluation.
        if area_ranges == [(None, None)]:
            fp[...] = 1
        else:
            det_areas = (
                det_bboxes[:, 2] - det_bboxes[:, 0] + extra_length) * (
                    det_bboxes[:, 3] - det_bboxes[:, 1] + extra_length)
            for i, (min_area, max_area) in enumerate(area_ranges):
                fp[i, (det_areas >= min_area) & (det_areas < max_area)] = 1

    if ioas is None or ioas.shape[1] <= 0:
        return tp, fp, det_bboxes
    else:
        # The evaluation of group-of TP and FP are done in two stages:
        # 1. All detections are first matched to non group-of boxes; true
        #    positives are determined.
        # 2. Detections that are determined as false positives are matched
        #    against group-of boxes and calculated group-of TP and FP.
        # Only used in OpenImages evaluation.
        det_bboxes_group = np.zeros(
            (num_scales, ioas.shape[1], det_bboxes.shape[1]), dtype=float)
        match_group_of = np.zeros((num_scales, num_dets), dtype=bool)
        tp_group = np.zeros((num_scales, num_gts_group), dtype=np.float32)
        ioas_max = ioas.max(axis=1)
        # for each det, which gt overlaps most with it
        ioas_argmax = ioas.argmax(axis=1)
        # sort all dets in descending order by scores
        sort_inds = np.argsort(-det_bboxes[:, -1])
        for k, (min_area, max_area) in enumerate(area_ranges):
            box_is_covered = tp[k]
            # if no area range is specified, gt_area_ignore is all False
            if min_area is None:
                gt_area_ignore = np.zeros_like(gt_ignore_inds, dtype=bool)
            else:
                gt_areas = (gt_bboxes[:, 2] - gt_bboxes[:, 0]) * (
                    gt_bboxes[:, 3] - gt_bboxes[:, 1])
                gt_area_ignore = (gt_areas < min_area) | (gt_areas >= max_area)
            for i in sort_inds:
                matched_gt = ioas_argmax[i]
                if not box_is_covered[i]:
                    if ioas_max[i] >= ioa_thr:
                        if not (gt_ignore_inds[matched_gt]
                                or gt_area_ignore[matched_gt]):
                            if not tp_group[k, matched_gt]:
                                tp_group[k, matched_gt] = 1
                                match_group_of[k, i] = True
                            else:
                                match_group_of[k, i] = True

                            if det_bboxes_group[k, matched_gt, -1] < \
                                    det_bboxes[i, -1]:
                                det_bboxes_group[k, matched_gt] = \
                                    det_bboxes[i]

        fp_group = (tp_group <= 0).astype(float)
        tps = []
        fps = []
        # concatenate tp, fp, and det-boxes which not matched group of
        # gt boxes and tp_group, fp_group, and det_bboxes_group which
        # matched group of boxes respectively.
        for i in range(num_scales):
            tps.append(
                np.concatenate((tp[i][~match_group_of[i]], tp_group[i])))
            fps.append(
                np.concatenate((fp[i][~match_group_of[i]], fp_group[i])))
            det_bboxes = np.concatenate(
                (det_bboxes[~match_group_of[i]], det_bboxes_group[i]))

        tp = np.vstack(tps)
        fp = np.vstack(fps)
        return tp, fp, det_bboxes


def get_cls_results(det_results, annotations, class_id):
    """Get det results and gt information of a certain class."""

    cls_dets = [img_res[class_id] for img_res in det_results]
    cls_gts = []
    cls_gts_ignore = []
    for ann in annotations:
        gt_inds = ann['labels'] == class_id
        cls_gts.append(ann['bboxes'][gt_inds, :])

        if ann.get('labels_ignore', None) is not None:
            ignore_inds = ann['labels_ignore'] == class_id
            cls_gts_ignore.append(ann['bboxes_ignore'][ignore_inds, :])
        else:
            cls_gts_ignore.append(np.empty((0, 4), dtype=np.float32))

    return cls_dets, cls_gts, cls_gts_ignore


def get_cls_group_ofs(annotations, class_id):
    """Get `gt_group_of` of a certain class, which is used in Open Images."""

    gt_group_ofs = []
    for ann in annotations:
        gt_inds = ann['labels'] == class_id
        if ann.get('gt_is_group_ofs', None) is not None:
            gt_group_ofs.append(ann['gt_is_group_ofs'][gt_inds])
        else:
            gt_group_ofs.append(np.empty((0, 1), dtype=np.bool))

    return gt_group_ofs


def eval_map(det_results,
             annotations,
             scale_ranges=None,
             iou_thr=0.5,
             ioa_thr=None,
             dataset=None,
             class_names=None,
             logger=None,
             tpfp_fn=None,
             nproc=4,
             use_legacy_coordinate=False,
             use_group_of=False):
    """Evaluate mAP of a dataset."""
    assert len(det_results) == len(annotations)
    if not use_legacy_coordinate:
        extra_length = 0.
    else:
        extra_length = 1.

    num_imgs = len(det_results)
    num_scales = len(scale_ranges) if scale_ranges is not None else 1
    num_classes = len(det_results[0])  # positive class num
    area_ranges = ([(rg[0]**2, rg[1]**2) for rg in scale_ranges]
                   if scale_ranges is not None else None)

    # There is no need to use multi processes to process
    # when num_imgs = 1 .
    if num_imgs > 1:
        assert nproc > 0, 'nproc must be at least one.'
        nproc = min(nproc, num_imgs)
        pool = Pool(nproc)

    eval_results = []
    for i in range(num_classes):
        # get gt and det bboxes of this class
        cls_dets, cls_gts, cls_gts_ignore = get_cls_results(
            det_results, annotations, i)
        # choose proper function according to datasets to compute tp and fp
        if tpfp_fn is None:
            if dataset in ['det', 'vid']:
                tpfp_fn = tpfp_imagenet
            elif dataset in ['oid_challenge', 'oid_v6'] \
                    or use_group_of is True:
                tpfp_fn = tpfp_openimages
            else:
                tpfp_fn = tpfp_default
        if not callable(tpfp_fn):
            raise ValueError(
                f'tpfp_fn has to be a function or None, but got {tpfp_fn}')

        if num_imgs > 1:
            # compute tp and fp for each image with multiple processes
            args = []
            if use_group_of:
                # used in Open Images Dataset evaluation
                gt_group_ofs = get_cls_group_ofs(annotations, i)
                args.append(gt_group_ofs)
                args.append([use_group_of for _ in range(num_imgs)])
            if ioa_thr is not None:
                args.append([ioa_thr for _ in range(num_imgs)])

            tpfp = pool.starmap(
                tpfp_fn,
                zip(cls_dets, cls_gts, cls_gts_ignore,
                    [iou_thr for _ in range(num_imgs)],
                    [area_ranges for _ in range(num_imgs)],
                    [use_legacy_coordinate for _ in range(num_imgs)], *args))
        else:
            tpfp = tpfp_fn(
                cls_dets[0],
                cls_gts[0],
                cls_gts_ignore[0],
                iou_thr,
                area_ranges,
                use_legacy_coordinate,
                gt_bboxes_group_of=(get_cls_group_ofs(annotations, i)[0]
                                    if use_group_of else None),
                use_group_of=use_group_of,
                ioa_thr=ioa_thr)
            tpfp = [tpfp]

        if use_group_of:
            tp, fp, cls_dets = tuple(zip(*tpfp))
        else:
            tp, fp = tuple(zip(*tpfp))
        # calculate gt number of each scale
        # ignored gts or gts beyond the specific scale are not counted
        num_gts = np.zeros(num_scales, dtype=int)
        for j, bbox in enumerate(cls_gts):
            if area_ranges is None:
                num_gts[0] += bbox.shape[0]
            else:
                gt_areas = (bbox[:, 2] - bbox[:, 0] + extra_length) * (
                    bbox[:, 3] - bbox[:, 1] + extra_length)
                for k, (min_area, max_area) in enumerate(area_ranges):
                    num_gts[k] += np.sum((gt_areas >= min_area)
                                         & (gt_areas < max_area))
        # sort all det bboxes by score, also sort tp and fp
        cls_dets = np.vstack(cls_dets)
        num_dets = cls_dets.shape[0]
        sort_inds = np.argsort(-cls_dets[:, -1])
        tp = np.hstack(tp)[:, sort_inds]
        fp = np.hstack(fp)[:, sort_inds]
        # calculate recall and precision with tp and fp
        tp = np.cumsum(tp, axis=1)
        fp = np.cumsum(fp, axis=1)
        eps = np.finfo(np.float32).eps
        recalls = tp / np.maximum(num_gts[:, np.newaxis], eps)
        precisions = tp / np.maximum((tp + fp), eps)
        # calculate AP
        if scale_ranges is None:
            recalls = recalls[0, :]
            precisions = precisions[0, :]
            num_gts = num_gts.item()
        mode = 'area' if dataset != 'voc07' else '11points'
        ap = average_precision(recalls, precisions, mode)
        eval_results.append({
            'num_gts': num_gts,
            'num_dets': num_dets,
            'recall': recalls,
            'precision': precisions,
            'ap': ap
        })

    if num_imgs > 1:
        pool.close()

    if scale_ranges is not None:
        # shape (num_classes, num_scales)
        all_ap = np.vstack([cls_result['ap'] for cls_result in eval_results])
        all_num_gts = np.vstack(
            [cls_result['num_gts'] for cls_result in eval_results])
        mean_ap = []
        for i in range(num_scales):
            if np.any(all_num_gts[:, i] > 0):
                mean_ap.append(all_ap[all_num_gts[:, i] > 0, i].mean())
            else:
                mean_ap.append(0.0)
    else:
        aps = []
        for cls_result in eval_results:
            if cls_result['num_gts'] > 0:
                aps.append(cls_result['ap'])
        mean_ap = np.array(aps).mean().item() if aps else 0.0

    print_map_summary(
        mean_ap, eval_results, class_names, area_ranges, logger=logger)
    return mean_ap, eval_results


def print_map_summary(mean_ap,
                      results,
                      class_names=None,
                      scale_ranges=None,
                      logger=None):
    """Print mAP and results of each class.

    A table will be printed to show the gts/dets/recall/AP of each class and
    the mAP.
    Args:
        mean_ap (float): Calculated from `eval_map()`.
        results (list[dict]): Calculated from `eval_map()`.
        dataset (list[str] | str | None): Dataset name or dataset classes.
        scale_ranges (list[tuple] | None): Range of scales to be evaluated.
        logger (logging.Logger | str | None): The way to print the mAP
            summary. See `mmengine.logging.print_log()` for details.
            Defaults to None.
    """

    if logger == 'silent':
        return

    if isinstance(results[0]['ap'], np.ndarray):
        num_scales = len(results[0]['ap'])
    else:
        num_scales = 1

    if scale_ranges is not None:
        assert len(scale_ranges) == num_scales

    num_classes = len(results)

    recalls = np.zeros((num_scales, num_classes), dtype=np.float32)
    aps = np.zeros((num_scales, num_classes), dtype=np.float32)
    num_gts = np.zeros((num_scales, num_classes), dtype=int)
    for i, cls_result in enumerate(results):
        if cls_result['recall'].size > 0:
            recalls[:, i] = np.array(cls_result['recall'], ndmin=2)[:, -1]
        aps[:, i] = cls_result['ap']
        num_gts[:, i] = cls_result['num_gts']

    if class_names is None:
        label_names = [str(i) for i in range(num_classes)]
    else:
        label_names = class_names

    if not isinstance(mean_ap, list):
        mean_ap = [mean_ap]

    header = ['class', 'gts', 'dets', 'recall', 'ap']
    for i in range(num_scales):
        if scale_ranges is not None:
            print_log(f'Scale range {scale_ranges[i]}', logger=logger)
        table_data = [header]
        for j in range(num_classes):
            row_data = [
                label_names[j], num_gts[i, j], results[j]['num_dets'],
                f'{recalls[i, j]:.3f}', f'{aps[i, j]:.3f}'
            ]
            table_data.append(row_data)
        table_data.append(['mAP', '', '', '', f'{mean_ap[i]:.3f}'])
        table = AsciiTable(table_data)
        table.inner_footing_row_border = True
        print_log('\n' + table.table, logger=logger)


def _recalls(all_ious, proposal_nums, thrs):

    img_num = all_ious.shape[0]
    total_gt_num = sum([ious.shape[0] for ious in all_ious])

    _ious = np.zeros((proposal_nums.size, total_gt_num), dtype=np.float32)
    for k, proposal_num in enumerate(proposal_nums):
        tmp_ious = np.zeros(0)
        for i in range(img_num):
            ious = all_ious[i][:, :proposal_num].copy()
            gt_ious = np.zeros((ious.shape[0]))
            if ious.size == 0:
                tmp_ious = np.hstack((tmp_ious, gt_ious))
                continue
            for j in range(ious.shape[0]):
                gt_max_overlaps = ious.argmax(axis=1)
                max_ious = ious[np.arange(0, ious.shape[0]), gt_max_overlaps]
                gt_idx = max_ious.argmax()
                gt_ious[j] = max_ious[gt_idx]
                box_idx = gt_max_overlaps[gt_idx]
                ious[gt_idx, :] = -1
                ious[:, box_idx] = -1
            tmp_ious = np.hstack((tmp_ious, gt_ious))
        _ious[k, :] = tmp_ious

    _ious = np.fliplr(np.sort(_ious, axis=1))
    recalls = np.zeros((proposal_nums.size, thrs.size))
    for i, thr in enumerate(thrs):
        recalls[:, i] = (_ious >= thr).sum(axis=1) / float(total_gt_num)

    return recalls


def set_recall_param(proposal_nums, iou_thrs):
    """Check proposal_nums and iou_thrs and set correct format."""
    if isinstance(proposal_nums, Sequence):
        _proposal_nums = np.array(proposal_nums)
    elif isinstance(proposal_nums, int):
        _proposal_nums = np.array([proposal_nums])
    else:
        _proposal_nums = proposal_nums

    if iou_thrs is None:
        _iou_thrs = np.array([0.5])
    elif isinstance(iou_thrs, Sequence):
        _iou_thrs = np.array(iou_thrs)
    elif isinstance(iou_thrs, float):
        _iou_thrs = np.array([iou_thrs])
    else:
        _iou_thrs = iou_thrs

    return _proposal_nums, _iou_thrs


def eval_recalls(gts,
                 proposals,
                 proposal_nums=None,
                 iou_thrs=0.5,
                 logger=None,
                 use_legacy_coordinate=False):
    """Calculate recalls."""

    img_num = len(gts)
    assert img_num == len(proposals)
    proposal_nums, iou_thrs = set_recall_param(proposal_nums, iou_thrs)
    all_ious = []
    for i in range(img_num):
        if proposals[i].ndim == 2 and proposals[i].shape[1] == 5:
            scores = proposals[i][:, 4]
            sort_idx = np.argsort(scores)[::-1]
            img_proposal = proposals[i][sort_idx, :]
        else:
            img_proposal = proposals[i]
        prop_num = min(img_proposal.shape[0], proposal_nums[-1])
        if gts[i] is None or gts[i].shape[0] == 0:
            ious = np.zeros((0, img_proposal.shape[0]), dtype=np.float32)
        else:
            ious = bbox_overlaps(
                gts[i],
                img_proposal[:prop_num, :4],
                use_legacy_coordinate=use_legacy_coordinate)
        all_ious.append(ious)
    all_ious = np.array(all_ious)
    recalls = _recalls(all_ious, proposal_nums, iou_thrs)

    print_recall_summary(recalls, proposal_nums, iou_thrs, logger=logger)
    return recalls


def print_recall_summary(recalls,
                         proposal_nums,
                         iou_thrs,
                         row_idxs=None,
                         col_idxs=None,
                         logger=None):
    """Print recalls in a table."""
    proposal_nums = np.array(proposal_nums, dtype=np.int32)
    iou_thrs = np.array(iou_thrs)
    if row_idxs is None:
        row_idxs = np.arange(proposal_nums.size)
    if col_idxs is None:
        col_idxs = np.arange(iou_thrs.size)
    row_header = [''] + iou_thrs[col_idxs].tolist()
    table_data = [row_header]
    for i, num in enumerate(proposal_nums[row_idxs]):
        row = [f'{val:.3f}' for val in recalls[row_idxs[i], col_idxs].tolist()]
        row.insert(0, num)
        table_data.append(row)
    table = AsciiTable(table_data)
    print_log('\n' + table.table, logger=logger)

jittordet/engine/evaluator/utils.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/evaluation/metrics/voc_metric.py
# Copyright (c) OpenMMLab. All rights reserved.
from collections import OrderedDict
from typing import List, Optional, Sequence, Union

import numpy as np

from jittordet.structures import BaseDataElement
from ..register import EVALUATORS
from .base_evaluator import BaseEvaluator
from .utils import eval_map, eval_recalls


@EVALUATORS.register_module()
class VocEvaluator(BaseEvaluator):
    """Pascal VOC evaluation metric.

    Args:
        iou_thrs (float or List[float]): IoU threshold. Defaults to 0.5.
        scale_ranges (List[tuple], optional): Scale ranges for evaluating
            mAP. If not specified, all bounding boxes would be included in
            evaluation. Defaults to None.
        metric (str | list[str]): Metrics to be evaluated. Options are
            'mAP', 'recall'. If is list, the first setting in the list will
             be used to evaluate metric.
        proposal_nums (Sequence[int]): Proposal number used for evaluating
            recalls, such as recall@100, recall@1000.
            Default: (100, 300, 1000).
        eval_mode (str): 'area' or '11points', 'area' means calculating the
            area under precision-recall curve, '11points' means calculating
            the average precision of recalls at [0, 0.1, ..., 1].
            The PASCAL VOC2007 defaults to use '11points', while PASCAL
            VOC2012 defaults to use 'area'.
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
            'gpu'. Defaults to 'cpu'.
        prefix (str, optional): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
            If prefix is not provided in the argument, self.default_prefix
            will be used instead. Defaults to None.
    """

    default_prefix: Optional[str] = 'pascal_voc'

    def __init__(self,
                 iou_thrs: Union[float, List[float]] = 0.5,
                 scale_ranges: Optional[List[tuple]] = None,
                 metric: Union[str, List[str]] = 'mAP',
                 proposal_nums: Sequence[int] = (100, 300, 1000),
                 eval_mode: str = '11points') -> None:
        self.results = []
        self.iou_thrs = [iou_thrs] if isinstance(iou_thrs, float) \
            else iou_thrs
        self.scale_ranges = scale_ranges
        # voc evaluation metrics
        if not isinstance(metric, str):
            assert len(metric) == 1
            metric = metric[0]
        allowed_metrics = ['recall', 'mAP']
        if metric not in allowed_metrics:
            raise KeyError(
                f"metric should be one of 'recall', 'mAP', but got {metric}.")
        self.metric = metric
        self.proposal_nums = proposal_nums
        assert eval_mode in ['area', '11points'], \
            'Unrecognized mode, only "area" and "11points" are supported'
        self.eval_mode = eval_mode

    # TODO: data_batch is no longer needed, consider adjusting the
    #  parameter position
    def process(self, dataset, data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and predictions.

        The processed
        results should be stored in ``self.results``, which will be used to
        compute the metrics when all batches have been processed.
        Args:
            data_batch (dict): A batch of data from the dataloader.
            data_samples (Sequence[dict]): A batch of data samples that
                contain annotations and predictions.
        """
        for data_sample in data_samples:
            if isinstance(data_sample, BaseDataElement):
                data_sample = data_sample.to_dict()

            pred = data_sample['pred_instances']
            pred_bboxes = pred['bboxes'].numpy()
            pred_scores = pred['scores'].numpy()
            pred_labels = pred['labels'].numpy()

            dets = []
            for label in range(len(dataset.metainfo['classes'])):
                index = np.where(pred_labels == label)[0]
                pred_bbox_scores = np.hstack(
                    [pred_bboxes[index], pred_scores[index].reshape((-1, 1))])
                dets.append(pred_bbox_scores)

            sample_idx = data_sample['sample_idx']
            gt_info = dataset.data_list[sample_idx]
            gt_labels, gt_bboxes = [], []
            gt_labels_ignore, gt_bboxes_ignore = [], []
            for ann in gt_info['instances']:
                if ann['ignore_flag']:
                    gt_bboxes_ignore.append(ann['bbox'])
                    gt_labels_ignore.append(ann['bbox_label'])
                else:
                    gt_bboxes.append(ann['bbox'])
                    gt_labels.append(ann['bbox_label'])
            gt_bboxes = np.array(gt_bboxes, dtype=np.float32) \
                if gt_bboxes else np.zeros((0, 4), dtype=np.float32)
            gt_bboxes_ignore = np.array(gt_bboxes_ignore, dtype=np.float32) \
                if gt_bboxes_ignore else np.zeros((0, 4), dtype=np.float32)
            gt_labels = np.array(gt_labels, dtype=np.int64) \
                if gt_labels else np.zeros((0, ), dtype=np.int64)
            gt_labels_ignore = np.array(gt_labels_ignore, dtype=np.int64) \
                if gt_labels_ignore else np.zeros((0, ), dtype=np.int64)
            gt = dict(
                labels=gt_labels,
                bboxes=gt_bboxes,
                labels_ignore=gt_labels_ignore,
                bboxes_ignore=gt_bboxes_ignore)
            self.results.append((gt, dets))

    def compute_metrics(self, dataset, results: list, logger) -> dict:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.
        Returns:
            dict: The computed metrics. The keys are the names of the metrics,
            and the values are corresponding results.
        """
        gts, preds = zip(*results)
        eval_results = OrderedDict()
        if self.metric == 'mAP':
            assert isinstance(self.iou_thrs, list)
            if self.eval_mode == '11points':
                dataset_name = 'voc07'
            else:
                dataset_name = 'voc12'
            class_names = dataset.metainfo['classes']

            mean_aps = []
            for iou_thr in self.iou_thrs:
                logger.info(f'\n{"-" * 15}iou_thr: {iou_thr}{"-" * 15}')
                # Follow the official implementation,
                # http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCdevkit_18-May-2011.tar
                # we should use the legacy coordinate system in mmdet 1.x,
                # which means w, h should be computed as 'x2 - x1 + 1` and
                # `y2 - y1 + 1`
                mean_ap, _ = eval_map(
                    preds,
                    gts,
                    scale_ranges=self.scale_ranges,
                    iou_thr=iou_thr,
                    dataset=dataset_name,
                    class_names=class_names,
                    logger=logger,
                    use_legacy_coordinate=True)
                mean_aps.append(mean_ap)
                eval_results[f'AP{int(iou_thr * 100):02d}'] = round(mean_ap, 3)
            eval_results['mAP'] = sum(mean_aps) / len(mean_aps)
            eval_results.move_to_end('mAP', last=False)
        elif self.metric == 'recall':
            # TODO: Currently not checked.
            gt_bboxes = [ann['bboxes'] for ann in self.annotations]
            recalls = eval_recalls(
                gt_bboxes,
                results,
                self.proposal_nums,
                self.iou_thrs,
                logger=logger,
                use_legacy_coordinate=True)
            for i, num in enumerate(self.proposal_nums):
                for j, iou_thr in enumerate(self.iou_thrs):
                    eval_results[f'recall@{num}@{iou_thr}'] = recalls[i, j]
            if recalls.shape[1] > 1:
                ar = recalls.mean(axis=1)
                for i, num in enumerate(self.proposal_nums):
                    eval_results[f'AR@{num}'] = ar[i]
        return eval_results

jittordet/engine/evaluator/voc_evaluator.py
JittorCV/jittordet
from .base_hook import BaseHook
from .checkpoint_hook import CheckpointHook
from .logger_hook import LoggerHook

__all__ = ['BaseHook', 'CheckpointHook', 'LoggerHook']

jittordet/engine/hooks/__init__.py
JittorCV/jittordet
# copy from OpenMMLab mmengine.hooks
# Copyright (c) OpenMMLab. All rights reserved.


class BaseHook:

    priority = 50

    def before_train(self, runner):
        pass

    def after_train(self, runner):
        pass

    def before_val(self, runner):
        pass

    def after_val(self, runner):
        pass

    def before_test(self, runner):
        pass

    def after_test(self, runner):
        pass

    def before_train_epoch(self, runner):
        self._before_epoch(runner, mode='train')

    def before_val_epoch(self, runner):
        self._before_epoch(runner, mode='val')

    def before_test_epoch(self, runner):
        self._before_epoch(runner, mode='test')

    def after_train_epoch(self, runner):
        self._after_epoch(runner, mode='train')

    def after_val_epoch(self, runner, metrics=None):
        self._after_epoch(runner, mode='val')

    def after_test_epoch(self, runner, metrics=None):
        self._after_epoch(runner, mode='test')

    def before_train_iter(self, runner, batch_idx, data_batch=None):
        self._before_iter(
            runner, batch_idx=batch_idx, data_batch=data_batch, mode='train')

    def before_val_iter(self, runner, batch_idx, data_batch=None):
        self._before_iter(
            runner, batch_idx=batch_idx, data_batch=data_batch, mode='val')

    def before_test_iter(self, runner, batch_idx, data_batch=None):
        self._before_iter(
            runner, batch_idx=batch_idx, data_batch=data_batch, mode='test')

    def after_train_iter(self,
                         runner,
                         batch_idx,
                         data_batch=None,
                         outputs=None):
        self._after_iter(
            runner,
            batch_idx=batch_idx,
            data_batch=data_batch,
            outputs=outputs,
            mode='train')

    def after_val_iter(self,
                       runner,
                       batch_idx,
                       data_batch=None,
                       outputs=None) -> None:
        self._after_iter(
            runner,
            batch_idx=batch_idx,
            data_batch=data_batch,
            outputs=outputs,
            mode='val')

    def after_test_iter(self,
                        runner,
                        batch_idx,
                        data_batch=None,
                        outputs=None) -> None:
        self._after_iter(
            runner,
            batch_idx=batch_idx,
            data_batch=data_batch,
            outputs=outputs,
            mode='test')

    def _before_epoch(self, runner, mode='train'):
        pass

    def _after_epoch(self, runner, mode='train'):
        pass

    def _before_iter(self,
                     runner,
                     batch_idx,
                     data_batch=None,
                     mode='train') -> None:
        pass

    def _after_iter(self,
                    runner,
                    batch_idx,
                    data_batch=None,
                    outputs=None,
                    mode='train'):
        pass

    def every_n_interval(self, idx, n):
        return (idx + 1) % n == 0 if n > 0 else False

jittordet/engine/hooks/base_hook.py
JittorCV/jittordet
import os.path as osp

from ..register import HOOKS
from .base_hook import BaseHook


@HOOKS.register_module()
class CheckpointHook(BaseHook):

    def __init__(self, interval=1, by_iter=False):
        assert isinstance(interval, (int, list))
        self.interval = interval
        self.by_iter = by_iter

    def after_train_epoch(self, runner):

        if self.by_iter:
            return None
        cur_epoch = runner.train_loop.cur_epoch

        if isinstance(self.interval, int):
            save_ckpt = self.every_n_interval(cur_epoch, self.interval)
        else:
            save_ckpt = cur_epoch in self.interval

        if save_ckpt:
            ckpt_filepath = osp.join(runner.log_dir,
                                     f'epoch_{cur_epoch + 1}.pkl')
            runner.save_checkpoint(ckpt_filepath)
            runner.logger.info(f'save checkpoint to {ckpt_filepath}')

    def after_train_iter(self,
                         runner,
                         batch_idx,
                         data_batch=None,
                         outputs=None):
        if not self.by_iter:
            return None
        cur_iter = runner.train_loop.cur_iter

        if isinstance(self.interval, int):
            save_ckpt = self.every_n_interval(cur_iter, self.interval)
        else:
            save_ckpt = cur_iter in self.interval

        if save_ckpt:
            ckpt_filepath = osp.join(runner.log_dir,
                                     f'iter_{cur_iter + 1}.pkl')
            runner.save_checkpoint(ckpt_filepath)
            runner.logger.info(f'save checkpoint to {ckpt_filepath}')

jittordet/engine/hooks/checkpoint_hook.py
JittorCV/jittordet
import datetime
import time
from collections import defaultdict

import jittor as jt
import numpy as np

from ..register import HOOKS
from .base_hook import BaseHook


@HOOKS.register_module()
class LoggerHook(BaseHook):

    def __init__(self,
                 interval=50,
                 interval_exp_name=300,
                 max_log_length=10000):
        self.interval = interval
        self.interval_exp_name = interval_exp_name
        self.max_log_length = max_log_length

        empty_ndarray = lambda: np.zeros([0])  # noqa E731
        self._train_log_history = defaultdict(empty_ndarray)
        self._val_log_history = defaultdict(empty_ndarray)
        self._test_log_history = defaultdict(empty_ndarray)

    def _select_log_history(self, phase):
        assert phase in ['train', 'val', 'test']
        if phase == 'train':
            log_history = self._train_log_history
        elif phase == 'val':
            log_history = self._val_log_history
        else:
            log_history = self._test_log_history
        return log_history

    def update_log_hitory(self, phase, key, value):
        # validate value
        if isinstance(value, jt.Var):
            if jt.in_mpi:
                value = value.mpi_all_reduce('mean')
            value = value.numpy()
        elif not isinstance(value, np.ndarray):
            value = np.array([value], dtype=np.float32)

        # record history
        log_history = self._select_log_history(phase)
        log_value = log_history[key]
        log_value = np.concatenate([log_value, value])
        if log_value.size > self.max_log_length:
            log_value = log_value[-self.max_log_length:]
        log_history[key] = log_value

    def get_log_hitory(self, phase, key, window_size=1, reduction='mean'):
        log_history = self._select_log_history(phase)
        log_value = log_history[key]
        if log_value.size == 0:
            raise RuntimeError(f'{key} has not been logged in history')

        if reduction == 'mean':
            window_size = min(window_size, log_value.size)
            log_value = log_value[-window_size:]
            log_value = log_value.mean().item()
        elif reduction == 'current':
            log_value = log_value[-1].item()
        else:
            raise KeyError('Only support "mean" and "current" reduction, but '
                           f'got {reduction}')
        return log_value

    def format_train_log_str(self, runner, batch_idx):
        log_str_list = []

        # epoch iteration number information
        log_str = '(train) '
        cur_epoch = runner.train_loop.cur_epoch
        max_epoch = runner.train_loop.max_epoch
        iter_per_epoch = len(runner.train_dataset)
        # get the max length of iteration and epoch number
        epoch_len = len(str(max_epoch))
        iter_len = len(str(iter_per_epoch))
        # right just the length
        cur_epoch = str(cur_epoch + 1).rjust(epoch_len)
        cur_iter = str(batch_idx + 1).rjust(iter_len)
        log_str += f'[{cur_epoch}/{max_epoch}]'
        log_str += f'[{cur_iter}/{iter_per_epoch}]'
        log_str_list.append(log_str)

        # iter time and etc time
        iter_time = self.get_log_hitory(
            'train', 'time', 1000, reduction='mean')
        past_iter = runner.train_loop.cur_iter
        total_iter = len(runner.train_dataset) * max_epoch
        eta_time = iter_time * (total_iter - past_iter)
        eta_time = datetime.timedelta(seconds=int(eta_time))
        mm, ss = divmod(eta_time.seconds, 60)
        hh, mm = divmod(mm, 60)
        format_eta_time = f'{eta_time.days:01d} day {hh:02d}:{mm:02d}:{ss:02d}'
        log_str_list.extend(
            [f'eta: {format_eta_time}', f'time: {iter_time:.4f}'])

        # leanring rate information
        lr = runner.optimizer.lr
        log_str_list.append(f'lr: {lr:.4e}')

        # other information
        for key in self._train_log_history.keys():
            if key == 'time':
                continue
            if 'loss' in key or 'acc' in key:
                log_value = self.get_log_hitory(
                    'train', key, self.interval, reduction='mean')
            else:
                log_value = self.get_log_hitory(
                    'train', key, self.interval, reduction='current')
            log_str_list.append(f'{key}: {log_value:.4f}')

        log_str = '  '.join(log_str_list)
        return log_str

    def format_test_val_log_str(self, runner, batch_idx, phase):
        assert phase in ['val', 'test']
        dataset = runner.val_dataset if phase == 'val' else runner.test_dataset
        loop = runner.val_loop if phase == 'val' else runner.test_loop
        log_str_list = []

        # epoch iteration number information
        log_str = '(val) [1/1]' if phase == 'val' else '(test) [1/1]'
        iter_per_epoch = len(dataset)
        # get the max length of iteration and epoch number
        iter_len = len(str(iter_per_epoch))
        # right just the length
        cur_iter = str(batch_idx + 1).rjust(iter_len)
        log_str += f'[{cur_iter}/{iter_per_epoch}]'
        log_str_list.append(log_str)

        # iter time and etc time
        iter_time = self.get_log_hitory(phase, 'time', 1000, reduction='mean')
        past_iter = loop.cur_iter
        eta_time = iter_time * (len(dataset) - past_iter)
        eta_time = datetime.timedelta(seconds=int(eta_time))
        mm, ss = divmod(eta_time.seconds, 60)
        hh, mm = divmod(mm, 60)
        format_eta_time = f'{eta_time.days:01d} day {hh:02d}:{mm:02d}:{ss:02d}'
        log_str_list.extend(
            [f'eta: {format_eta_time}', f'time: {iter_time:.4f}'])

        # other information
        for key in self._select_log_history(phase).keys():
            if key == 'time':
                continue
            log_value = self.get_log_hitory(
                phase, key, self.interval, reduction='current')
            log_str_list.append(f'{key}: {log_value:.4f}')

        log_str = '  '.join(log_str_list)
        return log_str

    def before_train_iter(self, runner, batch_idx, data_batch=None):
        self.train_iter_start_time = time.time()

    def before_val_iter(self, runner, batch_idx, data_batch=None):
        self.val_iter_start_time = time.time()

    def before_test_iter(self, runner, batch_idx, data_batch=None):
        self.test_iter_start_time = time.time()

    def after_train_iter(self,
                         runner,
                         batch_idx,
                         data_batch=None,
                         outputs=None):
        for key, value in outputs.items():
            self.update_log_hitory('train', key, value)
        iter_time = time.time() - self.train_iter_start_time
        self.update_log_hitory('train', 'time', iter_time)

        if self.every_n_interval(batch_idx, self.interval):
            log_str = self.format_train_log_str(runner, batch_idx)
            runner.logger.info(log_str)

        if self.every_n_interval(batch_idx, self.interval_exp_name):
            runner.logger.info(f'experiment name: {runner.experiment_name}')

    def after_val_iter(self, runner, batch_idx, data_batch=None, outputs=None):
        iter_time = time.time() - self.val_iter_start_time
        self.update_log_hitory('val', 'time', iter_time)

        if self.every_n_interval(batch_idx, self.interval):
            log_str = self.format_test_val_log_str(runner, batch_idx, 'val')
            runner.logger.info(log_str)

        if self.every_n_interval(batch_idx, self.interval_exp_name):
            runner.logger.info(f'experiment name: {runner.experiment_name}')

    def after_test_iter(self,
                        runner,
                        batch_idx,
                        data_batch=None,
                        outputs=None):
        iter_time = time.time() - self.test_iter_start_time
        self.update_log_hitory('test', 'time', iter_time)

        if self.every_n_interval(batch_idx, self.interval):
            log_str = self.format_test_val_log_str(runner, batch_idx, 'test')
            runner.logger.info(log_str)

        if self.every_n_interval(batch_idx, self.interval_exp_name):
            runner.logger.info(f'experiment name: {runner.experiment_name}')

    def after_val_epoch(self, runner, metrics=None):
        log_str = self.format_metrics(metrics)
        if log_str is not None:
            runner.logger.info(log_str)

    def after_test_epoch(self, runner, metrics=None):
        log_str = self.format_metrics(metrics)
        if log_str is not None:
            runner.logger.info(log_str)

    def format_metrics(self, metrics):
        log_str_list = []
        if isinstance(metrics, dict):
            for key, value in metrics.items():
                log_str_list.append(f'{key}: {value:.3f}')
        else:
            log_str_list.append(f'metrics: {metrics}')

        log_str = '  '.join(log_str_list) if log_str_list else None
        return log_str

jittordet/engine/hooks/logger_hook.py
JittorCV/jittordet
"""Modified from mmcv.utils.logging."""
import logging

import jittor as jt

__all__ = ['get_logger', 'print_log']
initialized_loggers = set()


def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'):
    logger = logging.getLogger(name)
    if name in initialized_loggers:
        return logger

    stream_handler = logging.StreamHandler()
    handlers = [stream_handler]

    if jt.rank == 0 and log_file is not None:
        file_handler = logging.FileHandler(log_file, file_mode)
        handlers.append(file_handler)

    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    for handler in handlers:
        handler.setFormatter(formatter)
        handler.setLevel(log_level)
        logger.addHandler(handler)

    if jt.rank == 0:
        logger.setLevel(log_level)
    else:
        logger.setLevel(logging.ERROR)

    initialized_loggers.add(name)

    return logger


def print_log(msg, logger=None, level=logging.INFO):
    if logger is None:
        print(msg)
    elif isinstance(logger, logging.Logger):
        logger.log(level, msg)
    elif logger == 'silent':
        pass
    elif isinstance(logger, str):
        _logger = get_logger(logger)
        _logger.log(level, msg)
    else:
        raise TypeError(
            'logger should be either a logging.Logger object, str, '
            f'"silent" or None, but got {type(logger)}')

jittordet/engine/logger.py
JittorCV/jittordet
from .base_loop import BaseLoop
from .test_loop import TestLoop
from .train_loop import EpochTrainLoop
from .val_loop import ValLoop

__all__ = ['BaseLoop', 'EpochTrainLoop', 'ValLoop', 'TestLoop']

jittordet/engine/loops/__init__.py
JittorCV/jittordet
from abc import ABCMeta, abstractmethod


class BaseLoop(metaclass=ABCMeta):

    def __init__(self, runner):
        self._runner = runner
        self._max_epoch = 1
        self._epoch = 0
        self._iter = 0

    @property
    def runner(self):
        return self._runner

    @property
    def cur_epoch(self):
        return self._epoch

    @property
    def max_epoch(self):
        return self._max_epoch

    @property
    def cur_iter(self):
        return self._iter

    @abstractmethod
    def run(self):
        pass

    def state_dict(self):
        return dict(epoch=self._epoch, iter=self._iter)

    def load_state_dict(self, data):
        assert isinstance(data, dict)
        self._epoch = data['epoch']
        self._iter = data['iter']

jittordet/engine/loops/base_loop.py
JittorCV/jittordet
import jittor as jt

from ..register import LOOPS
from .base_loop import BaseLoop


@LOOPS.register_module()
class TestLoop(BaseLoop):

    def run(self):
        """Launch validation."""
        self.runner.call_hook('before_test')
        self.runner.call_hook('before_test_epoch')
        self.runner.model.eval()
        self._iter = 0

        for idx, data_batch in enumerate(self.runner.test_dataset):
            self.run_iter(idx, data_batch)

        # compute metrics
        metrics = self.runner.test_evaluator.evaluate(self.runner.test_dataset,
                                                      self.runner.logger)

        self.runner.call_hook('after_test_epoch', metrics=metrics)
        self.runner.call_hook('after_test')

    @jt.no_grad()
    def run_iter(self, idx, data_batch):
        self.runner.call_hook(
            'before_test_iter', batch_idx=idx, data_batch=data_batch)

        outputs = self.runner.model(data_batch, phase='predict')
        self.runner.test_evaluator.process(self.runner.test_dataset, outputs)

        self.runner.call_hook(
            'after_test_iter',
            batch_idx=idx,
            data_batch=data_batch,
            outputs=outputs)
        self._iter += 1

jittordet/engine/loops/test_loop.py
JittorCV/jittordet
from ..register import LOOPS
from .base_loop import BaseLoop


@LOOPS.register_module()
class EpochTrainLoop(BaseLoop):

    def __init__(self, runner, max_epoch, val_interval=1):
        super().__init__(runner=runner)
        self.val_interval = val_interval
        self._max_epoch = max_epoch

    def run(self):
        self.runner.call_hook('before_train')

        while self._epoch < self._max_epoch:
            self.run_epoch()

            if self._epoch % self.val_interval == 0:
                if self.runner.val_loop is not None:
                    self.runner.val_loop.run()

        self.runner.call_hook('after_train')

    def run_epoch(self):
        """Iterate one epoch."""
        self.runner.call_hook('before_train_epoch')

        self.runner.model.train()
        for idx, data_batch in enumerate(self.runner.train_dataset):
            self.run_iter(idx, data_batch)

        for _scheduler in self.runner.scheduler:
            if not getattr(_scheduler, 'by_iter', False):
                _scheduler.step()

        self.runner.call_hook('after_train_epoch')
        self._epoch += 1

    def run_iter(self, idx, data_batch):
        """Iterate one min-batch."""
        self.runner.call_hook(
            'before_train_iter', batch_idx=idx, data_batch=data_batch)

        loss, loss_vars = self.runner.model(data_batch, phase='loss')
        self.runner.optimizer.step(loss)
        for _scheduler in self.runner.scheduler:
            # for warmup scheduler
            if getattr(_scheduler, 'by_iter', False):
                _scheduler.step()

        self.runner.call_hook(
            'after_train_iter',
            batch_idx=idx,
            data_batch=data_batch,
            outputs=loss_vars)
        self._iter += 1

jittordet/engine/loops/train_loop.py
JittorCV/jittordet
import jittor as jt

from ..register import LOOPS
from .base_loop import BaseLoop


@LOOPS.register_module()
class ValLoop(BaseLoop):

    def run(self):
        """Launch validation."""
        self.runner.call_hook('before_val')
        self.runner.call_hook('before_val_epoch')
        self.runner.model.eval()
        self._iter = 0

        for idx, data_batch in enumerate(self.runner.val_dataset):
            self.run_iter(idx, data_batch)

        # compute metrics
        metrics = self.runner.val_evaluator.evaluate(self.runner.val_dataset,
                                                     self.runner.logger)

        self.runner.call_hook('after_val_epoch', metrics=metrics)
        self.runner.call_hook('after_val')

    @jt.no_grad()
    def run_iter(self, idx, data_batch):
        self.runner.call_hook(
            'before_val_iter', batch_idx=idx, data_batch=data_batch)

        outputs = self.runner.model(data_batch, phase='predict')
        self.runner.val_evaluator.process(self.runner.val_dataset, outputs)

        self.runner.call_hook(
            'after_val_iter',
            batch_idx=idx,
            data_batch=data_batch,
            outputs=outputs)
        self._iter += 1

jittordet/engine/loops/val_loop.py
JittorCV/jittordet
from .optimizers import register_jittor_optim
from .schedulers import (BaseScheduler, CosineAnnealingLR, ExponentialLR,
                         StepLR, WarmUpLR)

__all__ = [
    'register_jittor_optim', 'BaseScheduler', 'WarmUpLR', 'CosineAnnealingLR',
    'ExponentialLR', 'StepLR'
]

jittordet/engine/optim/__init__.py
JittorCV/jittordet
import inspect

import jittor.optim as optim

from ..register import OPTIMIZERS


def register_jittor_optim():
    for name, module in optim.__dict__.items():
        if not inspect.isclass(module):
            continue

        if issubclass(module,
                      optim.Optimizer) and module is not optim.Optimizer:
            OPTIMIZERS.register_module(name=name, module=module)


register_jittor_optim()

jittordet/engine/optim/optimizers.py
JittorCV/jittordet
import jittor.lr_scheduler as scheduler

from ..register import SCHEDULERS


@SCHEDULERS.register_module()
class BaseScheduler:

    def state_dict(self):
        state_dict = {}
        exclude = ['optimizer']
        for key, value in self.__dict__.items():
            if key in exclude or callable(value):
                continue
            state_dict[key] = value
        return state_dict

    def load_state_dict(self, data):
        assert isinstance(data, dict)
        for key, value in data.items():
            if key in self.__dict__:
                self.__dict__[key] = value


@SCHEDULERS.register_module()
class WarmUpLR(BaseScheduler):
    """Copy from JDet. Warm LR scheduler, which is the base lr_scheduler,
    default we use it.

    Args:
        optimizer (Optimizer): the optimizer to optimize the model
        warmup (string): Type of warmup used. It can be None(use no warmup),
            'constant', 'linear' or 'exp'
        warmup_iters (int): The number of iterations or epochs that warmup
            lasts
        warmup_ratio (float): LR used at the beginning of warmup equals to
            warmup_ratio * initial_lr
    """

    def __init__(self,
                 optimizer,
                 warmup_ratio=1.0 / 3,
                 warmup_iters=500,
                 warmup='linear',
                 last_iter=-1):
        self.optimizer = optimizer
        self.warmup_ratio = warmup_ratio
        self.warmup_iters = warmup_iters
        self.warmup = warmup
        self.base_lr = optimizer.lr
        self.base_lr_pg = [
            pg.get('lr', optimizer.lr) for pg in optimizer.param_groups
        ]
        self.by_iter = True
        self.last_iter = last_iter

    def get_warmup_lr(self, lr, cur_iters):
        if self.warmup == 'constant':
            k = self.warmup_ratio
        elif self.warmup == 'linear':
            k = 1 - (1 - cur_iters / self.warmup_iters) * (1 -
                                                           self.warmup_ratio)
        elif self.warmup == 'exp':
            k = self.warmup_ratio**(1 - cur_iters / self.warmup_iters)
        return k * lr

    def _update_lr(self, steps):
        self.optimizer.lr = self.get_warmup_lr(self.base_lr, steps)
        for i, param_group in enumerate(self.optimizer.param_groups):
            param_group['lr'] = self.get_warmup_lr(self.base_lr_pg[i], steps)

    def step(self):
        self.last_iter += 1
        if self.last_iter <= self.warmup_iters:
            self._update_lr(self.last_iter)


@SCHEDULERS.register_module()
class CosineAnnealingLR(scheduler.CosineAnnealingLR, BaseScheduler):
    """CosineAnnealing LR Scheduler."""


@SCHEDULERS.register_module()
class ExponentialLR(scheduler.ExponentialLR, BaseScheduler):
    """Exponential LR Scheduler."""


@SCHEDULERS.register_module()
class StepLR(scheduler.StepLR, BaseScheduler):
    """Step LR Scheduler."""


@SCHEDULERS.register_module()
class MultiStepLR(scheduler.MultiStepLR, BaseScheduler):
    """Multiple Step LR Scheduler."""

jittordet/engine/optim/schedulers.py
JittorCV/jittordet
from .fields import (BATCH_SAMPLERS, BRICKS, DATASETS, EVALUATORS, HOOKS,
                     LOOPS, MODELS, OPTIMIZERS, SCHEDULERS, TASK_UTILS,
                     TRANSFORMS)
from .register import Register

__all__ = [
    'Register', 'LOOPS', 'HOOKS', 'OPTIMIZERS', 'SCHEDULERS', 'DATASETS',
    'MODELS', 'EVALUATORS', 'TRANSFORMS', 'BATCH_SAMPLERS', 'BRICKS',
    'TASK_UTILS'
]

jittordet/engine/register/__init__.py
JittorCV/jittordet
from .register import Register

# engine
LOOPS = Register('loops')
HOOKS = Register('hooks')
OPTIMIZERS = Register('optimizers')
SCHEDULERS = Register('schedulers')
EVALUATORS = Register('evaluators')

# dataset
DATASETS = Register('datasets')
TRANSFORMS = Register('transforms')
BATCH_SAMPLERS = Register('batch_sampler')

# model
MODELS = Register('models')
BRICKS = Register('bricks')
TASK_UTILS = Register('task_utils')

jittordet/engine/register/fields.py
JittorCV/jittordet
# modified from openmmlab mmengine
# https://github.com/open-mmlab/mmengine/blob/main/mmengine/registry/registry.py

import copy
import inspect
from functools import partial


class Register:

    def __init__(self, name):
        self._name = name
        self._module_dict = dict()

    def __len__(self):
        return len(self._module_dict)

    def __contains__(self, key):
        return key in self._module_dict

    def __repr__(self):
        format_str = self.__class__.__name__ + \
                     f'(name={self._name}, ' \
                     f'items={self._module_dict})'
        return format_str

    @property
    def name(self):
        return self._name

    @property
    def module_dict(self):
        return self._module_dict

    def get(self, key):
        return self._module_dict.get(key)

    def _register_module(self, module, module_name=None, force=False):
        if not inspect.isclass(module) and not inspect.isfunction(module):
            raise TypeError('module must be a class or a function, '
                            f'but got {type(module)}')

        if module_name is None:
            module_name = module.__name__
        if isinstance(module_name, str):
            module_name = [module_name]
        for name in module_name:
            if not force and name in self._module_dict:
                existed_module = self.module_dict[name]
                raise KeyError(f'{name} is already registered in {self.name} '
                               f'at {existed_module.__module__}')
            self._module_dict[name] = module

    def register_module(self, name=None, force=False, module=None):
        if not isinstance(force, bool):
            raise TypeError(f'force must be a boolean, but got {type(force)}')

        # raise the error ahead of time
        if not (name is None or isinstance(name, str)
                or isinstance(name, (list, tuple))):
            raise TypeError(
                'name must be None, an instance of str, or a sequence of str, '
                f'but got {type(name)}')

        # use it as a normal method: x.register_module(module=SomeClass)
        if module is not None:
            self._register_module(module=module, module_name=name, force=force)
            return module

        # use it as a decorator: @x.register_module()
        def _register(module):
            self._register_module(module=module, module_name=name, force=force)
            return module

        return _register

    def build(self, cfg, **default_args):
        if not isinstance(cfg, dict):
            raise TypeError(f'cfg should be a dict, but got {type(cfg)}')

        if not (isinstance(default_args, dict) or default_args is None):
            raise TypeError(
                f'default_args should be a dict, but got {type(default_args)}')

        if 'type' not in cfg:
            if default_args is None or 'type' not in default_args:
                raise KeyError(
                    '`cfg` or `default_args` must contain the key "type", '
                    f'but got {cfg}\n{default_args}')

        args = copy.deepcopy(cfg)
        if default_args is not None:
            for name, value in default_args.items():
                args.setdefault(name, value)

        obj_type = args.pop('type')
        if isinstance(obj_type, str):
            obj_cls = self.get(obj_type)
            if obj_cls is None:
                raise KeyError(
                    f'{obj_type} is not in the {self.name} registry.')
        elif inspect.isclass(obj_type) or inspect.isfunction(obj_type):
            obj_cls = obj_type
        else:
            raise TypeError(
                f'type must be a str or valid type, but got {type(obj_type)}')

        if inspect.isclass(obj_cls):
            if hasattr(obj_cls, 'from_cfg'):
                return obj_cls.from_cfg(*args)
            else:
                return obj_cls(**args)
        elif inspect.isfunction(obj_cls):
            return partial(obj_cls, **args)

jittordet/engine/register/register.py
JittorCV/jittordet
import copy
import inspect
import os
import os.path as osp
import random
import time

import jittor as jt
import jittor.nn as nn
import numpy as np
from jittor.dataset import Dataset
from jittor.optim import Optimizer

from .config import dump_cfg
from .evaluator import BaseEvaluator
from .hooks import BaseHook
from .logger import get_logger
from .loops import BaseLoop
from .register import (DATASETS, EVALUATORS, HOOKS, LOOPS, MODELS, OPTIMIZERS,
                       SCHEDULERS)

__all__ = ['Runner']


class Runner:

    def __init__(self,
                 model,
                 work_dir,
                 train_dataset=None,
                 val_dataset=None,
                 test_dataset=None,
                 train_loop=None,
                 val_loop=None,
                 test_loop=None,
                 val_evaluator=None,
                 test_evaluator=None,
                 optimizer=None,
                 scheduler=None,
                 hooks=None,
                 seed=None,
                 disable_cuda=False,
                 resume_from=None,
                 load_from=None,
                 experiment_name=None,
                 cfg=None):
        # setup work dir
        self._work_dir = osp.abspath(work_dir)
        if not osp.exists(self._work_dir) and jt.rank == 0:
            os.makedirs(self._work_dir)

        # hold on original configs
        self.cfg = cfg
        if self.cfg is not None:
            self.cfg = copy.deepcopy(self.cfg)

        # set random seed and other randomness related setting
        self.setup_randomness(seed)

        timestamp = jt.array(int(time.time()), dtype=jt.int64)
        if jt.in_mpi:
            timestamp = timestamp.mpi_broadcast(root=0)
        self._time_stamp = time.strftime('%Y%m%d_%H%M%S',
                                         time.localtime(timestamp.item()))

        # setup experiment name
        if experiment_name is not None:
            self._experiment_name = f'{experiment_name}_{self._time_stamp}'
        elif self.cfg.filename is not None:
            filename_no_ext = osp.splitext(osp.basename(self.cfg.filename))[0]
            self._experiment_name = f'{filename_no_ext}_{self._time_stamp}'
        else:
            self._experiment_name = self._time_stamp

        # setup hooks
        self.setup_hooks(hooks)

        self.disable_cuda = disable_cuda
        self._load_from = load_from
        self._resume_from = resume_from
        self.logger = None

        # build model
        self.model = self.build_model(model)

        # hold configs related to the training for lazy initialization
        self.train_dataset = train_dataset
        self.train_loop = train_loop
        self.optimizer = optimizer
        self.scheduler = scheduler

        # hold configs related to the val for lazy initialization
        self.val_dataset = val_dataset
        self.val_loop = val_loop
        self.val_evaluator = val_evaluator

        # hold configs related to the test for lazy initialization
        self.test_dataset = test_dataset
        self.test_loop = test_loop
        self.test_evaluator = test_evaluator

    @classmethod
    def from_cfg(cls, cfg):
        cfg = copy.deepcopy(cfg)
        runner = cls(
            model=cfg['model'],
            work_dir=cfg['work_dir'],
            train_dataset=cfg.get('train_dataset'),
            val_dataset=cfg.get('val_dataset'),
            test_dataset=cfg.get('test_dataset'),
            train_loop=cfg.get('train_loop'),
            val_loop=cfg.get('val_loop'),
            test_loop=cfg.get('test_loop'),
            val_evaluator=cfg.get('val_evaluator'),
            test_evaluator=cfg.get('test_evaluator'),
            optimizer=cfg.get('optimizer'),
            scheduler=cfg.get('scheduler'),
            hooks=cfg.get('hooks'),
            seed=cfg.get('seed'),
            disable_cuda=cfg.get('disable_cuda', False),
            resume_from=cfg.get('resume_from'),
            load_from=cfg.get('load_from'),
            experiment_name=cfg.get('experiment_name'),
            cfg=cfg)
        return runner

    @property
    def work_dir(self):
        return self._work_dir

    @property
    def time_stamp(self):
        return self._time_stamp

    @property
    def experiment_name(self):
        return self._experiment_name

    def init_model_weights(self):
        self.logger.info('Start initialize model weights.')

        def dfs_run_init_weights(model):
            for m in model._modules.values():
                dfs_run_init_weights(m)

            if hasattr(model, 'init_weights'):
                model.init_weights()

        dfs_run_init_weights(self.model)
        self.logger.info('Initialize succeed.')

    def build_model(self, model):
        if isinstance(model, nn.Module):
            return model
        elif isinstance(model, dict):
            return MODELS.build(model)
        else:
            raise TypeError('model should be a nn.Module object or dict')

    def build_dataset(self, dataset):
        if isinstance(dataset, Dataset):
            return dataset
        elif isinstance(dataset, dict):

            return DATASETS.build(dataset)
        else:
            raise TypeError(
                'dataset should be a jittor.dataset.Dataset object or dict')

    def build_optimizer(self, optimizer):
        if isinstance(optimizer, Optimizer):
            return Optimizer

        assert isinstance(optimizer, dict)
        optimizer = copy.deepcopy(optimizer)
        cfg_clip_grad = optimizer.pop('clip_grad', None)
        optimizer = OPTIMIZERS.build(optimizer, params=self.model.parameters())
        # hold clip grad config
        optimizer._cfg_clip_grad = cfg_clip_grad

        return optimizer

    def build_scheduler(self, schedulers, optimizer):
        assert isinstance(optimizer, Optimizer), \
            'optimzer should be built before building scheduler.'

        if not isinstance(schedulers, (list, tuple)):
            schedulers = [schedulers]

        param_schedulers = []
        for scheduler in schedulers:
            # recently, scheduler doesn't have a base class
            if not inspect.isclass(scheduler):
                assert isinstance(scheduler, dict)
                scheduler = SCHEDULERS.build(scheduler, optimizer=optimizer)
            param_schedulers.append(scheduler)

        # setup scheduler last_step from -1 to 0
        for scheduler in param_schedulers:
            scheduler.step()
        return param_schedulers

    def build_loop(self, loop):
        if isinstance(loop, BaseLoop):
            return loop
        elif isinstance(loop, dict):
            return LOOPS.build(loop, runner=self)
        else:
            raise TypeError('loops should be a BaseLoop object or dict')

    def build_evaluator(self, evaluator):
        if isinstance(evaluator, BaseEvaluator):
            return evaluator
        elif isinstance(evaluator, dict):
            return EVALUATORS.build(evaluator)
        else:
            raise TypeError(
                'evaluator should be a BaseEvaluator object or dict')

    def train(self):
        self.setup_logger_dir('train')

        train_related = (self.train_dataset, self.train_loop, self.optimizer,
                         self.scheduler)
        assert all([item is not None for item in train_related]), \
            'In training phase, "train_dataset", "train_loop", "optimizer", ' \
            'and "scheduler" should not be None.'
        self.train_dataset = self.build_dataset(self.train_dataset)
        self.train_loop = self.build_loop(self.train_loop)
        self.optimizer = self.build_optimizer(self.optimizer)
        self.scheduler = self.build_scheduler(self.scheduler, self.optimizer)

        val_related = (self.val_dataset, self.val_loop, self.val_evaluator)
        if any([item is not None for item in val_related]):
            assert all([item is not None for item in val_related]), \
                'In training phase, "val_dataset", "val_loop", ' \
                'and "val_evaluator" should be all None or not None.'
            self.val_dataset = self.build_dataset(self.val_dataset)
            self.val_loop = self.build_loop(self.val_loop)
            self.val_evaluator = self.build_evaluator(self.val_evaluator)

        # initialization
        if not self.load_or_resume():
            self.init_model_weights()
        if not self.disable_cuda:
            jt.flags.use_cuda = 1

        # start run
        self.train_loop.run()

    def val(self):
        self.setup_logger_dir('val')

        val_related = (self.val_dataset, self.val_loop, self.val_evaluator)
        assert all([item is not None for item in val_related]), \
            'In val phase, "val_dataset", "val_loop", and "val_evaluator" ' \
            'should be all None or not None.'
        self.val_dataset = self.build_dataset(self.val_dataset)
        self.val_loop = self.build_loop(self.val_loop)
        self.val_evaluator = self.build_evaluator(self.val_evaluator)

        # initialization
        if not self.load_or_resume():
            raise RuntimeError('Model has not been load in val')
        if not self.disable_cuda:
            jt.flags.use_cuda = 1

        # start run
        self.val_loop.run()

    def test(self):
        self.setup_logger_dir('test')

        test_related = (self.test_dataset, self.test_loop, self.test_evaluator)
        assert all([item is not None for item in test_related]), \
            'In test phase, "test_dataset", "test_loop", and ' \
            '"test_evaluator" should be all None or not None.'
        self.test_dataset = self.build_dataset(self.test_dataset)
        self.test_loop = self.build_loop(self.test_loop)
        self.test_evaluator = self.build_evaluator(self.test_evaluator)

        # initialization
        if not self.load_or_resume():
            raise RuntimeError('Model has not been load in test')
        if not self.disable_cuda:
            jt.flags.use_cuda = 1

        # start run
        self.test_loop.run()

    def setup_logger_dir(self, phase):
        assert phase in ['train', 'val', 'test']
        log_dir = osp.join(self.work_dir, self.time_stamp + '_' + phase)
        log_file = osp.join(log_dir, self.time_stamp + '.log')
        if jt.rank == 0 and not osp.exists(log_dir):
            os.makedirs(log_dir)
        self.log_dir = log_dir
        self.logger = get_logger(name='jittordet', log_file=log_file)

        cfg_filename = self.cfg.get('filename', None)
        cfg_filename = 'config.yml' if cfg_filename is None else \
            osp.basename(cfg_filename)
        if jt.rank == 0:
            dump_cfg(self.cfg, osp.join(self.log_dir, cfg_filename))

    def setup_randomness(self, seed):
        if seed is not None:
            rank_diff_seed = seed + jt.rank
            random.seed(rank_diff_seed)
            np.random.seed(rank_diff_seed)
            jt.set_global_seed(seed)

    def load_or_resume(self):
        assert not (self._load_from and self._resume_from), \
            'Can only set either "load_from" or "resume_from"'
        if self._load_from:
            self.load(self._load_from)
            return True
        if self._resume_from:
            self.resume(self._resume_from)
            return True
        return False

    def resume(self, checkpoint):
        self.logger.info(f'Start resume state dict from {checkpoint}')
        data = jt.load(checkpoint)
        model_state_dict = data['state_dict']
        self.model.load_state_dict(model_state_dict)

        # load train_loop info
        if self.train_loop is not None and 'loop' in data:
            self.train_loop.load_state_dict(data['loop'])

        # load optimizer info
        if self.optimizer is not None and 'optimizer' in data:
            self.optimizer.load_state_dict(data['optimizer'])

        # load scheduler info
        if self.scheduler is not None and 'scheduler' in data:
            for scheduler in self.scheduler:
                scheduler_name = scheduler.__class__.__name__
                if scheduler_name in data['scheduler']:
                    scheduler.load_state_dict(
                        data['scheduler'][scheduler_name])
        self.logger.info('Finish resuming')

    def load(self, checkpoint):
        self.logger.info(f'Start load parameters from {checkpoint}')
        data = jt.load(checkpoint)
        model_state_dict = data['state_dict']
        self.model.load_state_dict(model_state_dict)
        self.logger.info('Finish loading checkpoint')

    @jt.single_process_scope()
    def save_checkpoint(self, filepath, end_epoch=True):
        data = dict()
        data['state_dict'] = self.model.state_dict()

        # train_loop info
        loop_state_dict = self.train_loop.state_dict()
        if end_epoch:
            loop_state_dict['epoch'] += 1
        data['loop'] = loop_state_dict

        # optimizer info
        data['optimizer'] = self.optimizer.state_dict()

        # scheduler info
        scheduler_state_dict = dict()
        for scheduler in self.scheduler:
            scheduler_name = scheduler.__class__.__name__
            scheduler_state_dict[scheduler_name] = scheduler.state_dict()
        data['scheduler'] = scheduler_state_dict
        jt.save(data, filepath)

    def setup_hooks(self, hooks):
        if hooks is None:
            hooks = []
        assert isinstance(hooks, list)

        default_hooks = [dict(type='LoggerHook'), dict(type='CheckpointHook')]
        hook_keys = [hook['type'] for hook in hooks]
        for hook in default_hooks:
            if hook['type'] not in hook_keys:
                hooks.append(hook)

        _hooks = []
        for hook in hooks:
            if not isinstance(hook, BaseHook):
                hook = HOOKS.build(hook)
            _hooks.append(hook)
        self._hooks = sorted(_hooks, key=lambda x: x.priority, reverse=True)

    def call_hook(self, func_name, **kwargs):
        for hook in self._hooks:
            getattr(hook, func_name)(self, **kwargs)

jittordet/engine/runner.py
JittorCV/jittordet
from .backbones import *  # noqa: F401, F403
from .dense_heads import *  # noqa: F401, F403
from .frameworks import *  # noqa: F401, F403
from .layers import *  # noqa: F401, F403
from .losses import *  # noqa: F401, F403
from .necks import *  # noqa: F401, F403
from .preprocessors import *  # noqa: F401, F403
from .roi_heads import *  # noqa: F401, F403
from .task_utils import *  # noqa: F401, F403
from .utils import *  # noqa: F401, F403

jittordet/models/__init__.py
JittorCV/jittordet
# Modified from jdet/python/models/backbones/resnet.py
import jittor as jt
from jittor import nn

from jittordet.engine import MODELS


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    conv = nn.Conv(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=dilation,
        groups=groups,
        bias=False,
        dilation=dilation)
    jt.init.relu_invariant_gauss_(conv.weight, mode='fan_out')
    return conv


def conv1x1(in_planes, out_planes, stride=1):
    conv = nn.Conv(
        in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
    jt.init.relu_invariant_gauss_(conv.weight, mode='fan_out')
    return conv


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 groups=1,
                 base_width=64,
                 dilation=1,
                 norm_layer=None):
        super(BasicBlock, self).__init__()
        if (norm_layer is None):
            norm_layer = nn.BatchNorm
        if ((groups != 1) or (base_width != 64)):
            raise ValueError(
                'BasicBlock only supports groups=1 and base_width=64')
        if (dilation > 1):
            raise NotImplementedError(
                'Dilation > 1 not supported in BasicBlock')
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.Relu()
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def execute(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if (self.downsample is not None):
            identity = self.downsample(x)
        out += identity
        out = self.relu(out)
        return out


class Bottleneck(nn.Module):

    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 groups=1,
                 base_width=64,
                 dilation=1,
                 norm_layer=None):
        super(Bottleneck, self).__init__()
        if (norm_layer is None):
            norm_layer = nn.BatchNorm
        width = (int((planes * (base_width / 64.0))) * groups)
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, (planes * self.expansion))
        self.bn3 = norm_layer((planes * self.expansion))
        self.relu = nn.Relu()
        self.downsample = downsample
        self.stride = stride

    def execute(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if (self.downsample is not None):
            identity = self.downsample(x)
        out += identity
        out = self.relu(out)
        return out


@MODELS.register_module()
class ResNet(nn.Module):

    structures = {
        18: {
            'layers': [2, 2, 2, 2],
            'block': BasicBlock
        },
        34: {
            'layers': [3, 4, 6, 3],
            'block': BasicBlock
        },
        50: {
            'layers': [3, 4, 6, 3],
            'block': Bottleneck
        },
        101: {
            'layers': [3, 4, 23, 3],
            'block': Bottleneck
        },
        152: {
            'layers': [3, 8, 36, 3],
            'block': Bottleneck
        },
    }

    def __init__(self,
                 depth,
                 return_stages=['layer4'],
                 frozen_stages=-1,
                 norm_eval=True,
                 num_classes=None,
                 groups=1,
                 width_per_group=64,
                 replace_stride_with_dilation=None,
                 pretrained=None,
                 norm_layer=None):
        super(ResNet, self).__init__()
        if (norm_layer is None):
            norm_layer = nn.BatchNorm
        self.frozen_stages = frozen_stages
        self.norm_eval = norm_eval
        self._norm_layer = norm_layer
        self.inplanes = 64
        self.dilation = 1
        if (replace_stride_with_dilation is None):
            replace_stride_with_dilation = [False, False, False]
        if (len(replace_stride_with_dilation) != 3):
            raise ValueError(
                'replace_stride_with_dilation should be None or a 3-element '
                f'tuple, but got {replace_stride_with_dilation}')
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv(
            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        jt.init.relu_invariant_gauss_(self.conv1.weight, mode='fan_out')
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.Relu()
        self.maxpool = nn.Pool(
            kernel_size=3, stride=2, padding=1, op='maximum')

        assert depth in self.structures
        structure = self.structures[depth]
        block, layers = structure['block'], structure['layers']
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(
            block,
            128,
            layers[1],
            stride=2,
            dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(
            block,
            256,
            layers[2],
            stride=2,
            dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(
            block,
            512,
            layers[3],
            stride=2,
            dilate=replace_stride_with_dilation[2])
        self.num_classes = num_classes
        self.return_stages = return_stages
        if num_classes is not None:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
            self.fc = nn.Linear((512 * block.expansion), num_classes)
        self.pretrained = pretrained
        self._freeze_stages()

    def init_weights(self):
        assert self.pretrained is not None
        self.load(self.pretrained)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if ((stride != 1) or (self.inplanes != (planes * block.expansion))):
            downsample = nn.Sequential(
                conv1x1(self.inplanes, (planes * block.expansion), stride),
                norm_layer((planes * block.expansion)))
        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, self.groups,
                  self.base_width, previous_dilation, norm_layer))
        self.inplanes = (planes * block.expansion)
        for _ in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    groups=self.groups,
                    base_width=self.base_width,
                    dilation=self.dilation,
                    norm_layer=norm_layer))
        return nn.Sequential(*layers)

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.bn1.eval()
            for m in [self.conv1, self.bn1]:
                for param in m.parameters():
                    param.stop_grad()

        for i in range(1, self.frozen_stages + 1):
            m = getattr(self, 'layer{}'.format(i))
            m.eval()
            for param in m.parameters():
                param.stop_grad()

    def execute(self, x):
        outputs = []
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        for i in range(1, 5):
            name = f'layer{i}'
            x = getattr(self, name)(x)
            if name in self.return_stages:
                outputs.append(x)
        if self.num_classes is not None:
            x = self.avgpool(x)
            x = jt.reshape(x, (x.shape[0], -1))
            x = self.fc(x)
            if 'fc' in self.return_stages:
                outputs.append(x)
        return tuple(outputs)

    def train(self):
        super(ResNet, self).train()
        self._freeze_stages()
        if self.norm_eval:
            for m in self.modules():
                # trick: eval have effect on BatchNorm only
                if isinstance(m, nn.BatchNorm):
                    m.eval()


@MODELS.register_module()
class ResNetV1d(nn.Module):

    structures = {
        18: {
            'layers': [2, 2, 2, 2],
            'block': BasicBlock
        },
        34: {
            'layers': [3, 4, 6, 3],
            'block': BasicBlock
        },
        50: {
            'layers': [3, 4, 6, 3],
            'block': Bottleneck
        },
        101: {
            'layers': [3, 4, 23, 3],
            'block': Bottleneck
        },
        152: {
            'layers': [3, 8, 36, 3],
            'block': Bottleneck
        },
    }

    def __init__(self,
                 depth,
                 return_stages=['layer4'],
                 frozen_stages=-1,
                 norm_eval=True,
                 num_classes=None,
                 groups=1,
                 width_per_group=64,
                 replace_stride_with_dilation=None,
                 norm_layer=None,
                 pretrained=None):
        super(ResNetV1d, self).__init__()
        if (norm_layer is None):
            norm_layer = nn.BatchNorm
        self.frozen_stages = frozen_stages
        self.norm_eval = norm_eval
        self._norm_layer = norm_layer
        self.inplanes = 64
        self.dilation = 1
        if (replace_stride_with_dilation is None):
            replace_stride_with_dilation = [False, False, False]
        if (len(replace_stride_with_dilation) != 3):
            raise ValueError(
                'replace_stride_with_dilation should be None or a 3-element '
                f'tuple, got {replace_stride_with_dilation}')
        self.groups = groups
        self.base_width = width_per_group
        self.C1 = nn.Sequential(
            nn.Conv(3, 32, kernel_size=3, stride=2, padding=1, bias=False),
            norm_layer(32),
            nn.Relu(),
            nn.Conv(32, 32, kernel_size=3, stride=1, padding=1, bias=False),
            norm_layer(32),
            nn.Relu(),
            nn.Conv(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            norm_layer(64),
            nn.Relu(),
        )

        assert depth in self.structures
        structure = self.structures[depth]
        block, layers = structure['block'], structure['layers']

        self.relu = nn.Relu()
        self.maxpool = nn.Pool(
            kernel_size=3, stride=2, padding=1, op='maximum')
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(
            block,
            128,
            layers[1],
            stride=2,
            dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(
            block,
            256,
            layers[2],
            stride=2,
            dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(
            block,
            512,
            layers[3],
            stride=2,
            dilate=replace_stride_with_dilation[2])
        self.num_classes = num_classes
        self.return_stages = return_stages
        if num_classes is not None:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
            self.fc = nn.Linear((512 * block.expansion), num_classes)

        self.pretrained = pretrained
        self._freeze_stages()

    def init_weights(self):
        assert self.pretrained is not None
        self.load(self.pretrained)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if ((stride != 1) or (self.inplanes != (planes * block.expansion))):
            downsample = nn.Sequential(
                nn.Pool(stride, stride=stride, op='mean'),
                conv1x1(self.inplanes, (planes * block.expansion), 1),
                norm_layer((planes * block.expansion)))
        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, self.groups,
                  self.base_width, previous_dilation, norm_layer))
        self.inplanes = (planes * block.expansion)
        for _ in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    groups=self.groups,
                    base_width=self.base_width,
                    dilation=self.dilation,
                    norm_layer=norm_layer))
        return nn.Sequential(*layers)

    def execute(self, x):
        outputs = []
        x = self.C1(x)
        x = self.maxpool(x)
        for i in range(1, 5):
            name = f'layer{i}'
            x = getattr(self, name)(x)
            if name in self.return_stages:
                outputs.append(x)
        if self.num_classes is not None:
            x = self.avgpool(x)
            x = jt.reshape(x, (x.shape[0], -1))
            x = self.fc(x)
            if 'fc' in self.return_stages:
                outputs.append(x)
        return outputs

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.bn1.eval()
            for m in [self.conv1, self.bn1]:
                for param in m.parameters():
                    param.stop_grad()

        for i in range(1, self.frozen_stages + 1):
            m = getattr(self, 'layer{}'.format(i))
            m.eval()
            for param in m.parameters():
                param.stop_grad()

    def train(self):
        super(ResNetV1d, self).train()
        self._freeze_stages()
        if self.norm_eval:
            for m in self.modules():
                # trick: eval have effect on BatchNorm only
                if isinstance(m, nn.BatchNorm):
                    m.eval()

jittordet/models/backbones/resnet.py
JittorCV/jittordet
from .anchor_head import AnchorHead
from .base_dense_head import BaseDenseHead
from .retina_head import RetinaHead
from .rpn_head import RPNHead

__all__ = ['BaseDenseHead', 'AnchorHead', 'RetinaHead', 'RPNHead']

jittordet/models/dense_heads/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/dense_heads/anchor_head.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS, TASK_UTILS, ConfigType, OptConfigType
from jittordet.structures import InstanceData, InstanceList, OptInstanceList
from ..task_utils import PseudoSampler, anchor_inside_flags
from ..utils import images_to_levels, multi_apply, normal_init, unmap
from .base_dense_head import BaseDenseHead


@MODELS.register_module()
class AnchorHead(BaseDenseHead):
    """Anchor-based head (RPN, RetinaNet, SSD, etc.).

    Args:
        num_classes (int): Number of categories excluding the background
            category.
        in_channels (int): Number of channels in the input feature map.
        feat_channels (int): Number of hidden channels. Used in child classes.
        anchor_generator (dict): Config dict for anchor generator
        bbox_coder (dict): Config of bounding box coder.
        reg_decoded_bbox (bool): If true, the regression loss would be
            applied directly on decoded bounding boxes, converting both
            the predicted boxes and regression targets to absolute
            coordinates format. Default False. It should be `True` when
            using `IoULoss`, `GIoULoss`, or `DIoULoss` in the bbox head.
        loss_cls (dict): Config of classification loss.
        loss_bbox (dict): Config of localization loss.
        train_cfg (dict): Training config of anchor head.
        test_cfg (dict): Testing config of anchor head.
        init_cfg (dict or list[dict], optional): Initialization config dict.
    """  # noqa: W605

    def __init__(
        self,
        num_classes: int,
        in_channels: int,
        feat_channels: int = 256,
        anchor_generator: ConfigType = dict(
            type='AnchorGenerator',
            scales=[8, 16, 32],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder: ConfigType = dict(
            type='DeltaXYWHBBoxCoder',
            clip_border=True,
            target_means=(.0, .0, .0, .0),
            target_stds=(1.0, 1.0, 1.0, 1.0)),
        reg_decoded_bbox: bool = False,
        loss_cls: ConfigType = dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox: ConfigType = dict(
            type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0),
        train_cfg: OptConfigType = None,
        test_cfg: OptConfigType = None,
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.feat_channels = feat_channels
        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
        if self.use_sigmoid_cls:
            self.cls_out_channels = num_classes
        else:
            self.cls_out_channels = num_classes + 1

        if self.cls_out_channels <= 0:
            raise ValueError(f'num_classes={num_classes} is too small')
        self.reg_decoded_bbox = reg_decoded_bbox

        self.bbox_coder = TASK_UTILS.build(bbox_coder)
        self.loss_cls = MODELS.build(loss_cls)
        self.loss_bbox = MODELS.build(loss_bbox)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        if self.train_cfg:
            self.assigner = TASK_UTILS.build(self.train_cfg['assigner'])
            if train_cfg.get('sampler', None) is not None:
                self.sampler = TASK_UTILS.build(
                    self.train_cfg['sampler'], default_args=dict(context=self))
            else:
                self.sampler = PseudoSampler(context=self)

        self.prior_generator = TASK_UTILS.build(anchor_generator)

        # Usually the numbers of anchors for each level are the same
        # except SSD detectors. So it is an int in the most dense
        # heads but a list of int in SSDHead
        self.num_base_priors = self.prior_generator.num_base_priors[0]
        self._init_layers()

    def _init_layers(self) -> None:
        """Initialize layers of the head."""
        self.conv_cls = nn.Conv2d(self.in_channels,
                                  self.num_base_priors * self.cls_out_channels,
                                  1)
        self.conv_reg = nn.Conv2d(self.in_channels, self.num_base_priors * 4,
                                  1)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.01)

    def execute_single(self, x: jt.Var) -> Tuple[jt.Var, jt.Var]:
        """Forward feature of a single scale level.

        Args:
            x (Tensor): Features of a single scale level.

        Returns:
            tuple:
                cls_score (Tensor): Cls scores for a single scale level \
                    the channels number is num_base_priors * num_classes.
                bbox_pred (Tensor): Box energies / deltas for a single scale \
                    level, the channels number is num_base_priors * 4.
        """
        cls_score = self.conv_cls(x)
        bbox_pred = self.conv_reg(x)
        return cls_score, bbox_pred

    def execute(self, x: Tuple[jt.Var]) -> Tuple[List[jt.Var]]:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and bbox prediction.

                - cls_scores (list[Tensor]): Classification scores for all \
                    scale levels, each is a 4D-tensor, the channels number \
                    is num_base_priors * num_classes.
                - bbox_preds (list[Tensor]): Box energies / deltas for all \
                    scale levels, each is a 4D-tensor, the channels number \
                    is num_base_priors * 4.
        """
        return multi_apply(self.execute_single, x)

    def get_priors(self,
                   featmap_sizes: List[tuple],
                   batch_img_metas: List[dict]) \
            -> Tuple[List[List[jt.Var]], List[List[jt.Var]]]:
        """Get anchors according to feature map sizes.

        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            batch_img_metas (list[dict]): Image meta info.
            device (torch.device | str): Device for returned tensors.
                Defaults to cuda.

        Returns:
            tuple:

                - anchor_list (list[list[Tensor]]): Anchors of each image.
                - valid_flag_list (list[list[Tensor]]): Valid flags of each
                  image.
        """
        num_imgs = len(batch_img_metas)

        # since feature map sizes of all images are the same, we only compute
        # anchors for one time
        multi_level_anchors = self.prior_generator.grid_priors(featmap_sizes)
        anchor_list = [multi_level_anchors for _ in range(num_imgs)]

        # for each image, we compute valid flags of multi level anchors
        valid_flag_list = []
        for img_id, img_meta in enumerate(batch_img_metas):
            multi_level_flags = self.prior_generator.valid_flags(
                featmap_sizes, img_meta['pad_shape'])
            valid_flag_list.append(multi_level_flags)

        return anchor_list, valid_flag_list

    def _get_targets_single(self,
                            flat_anchors: jt.Var,
                            valid_flags: jt.Var,
                            gt_instances: InstanceData,
                            img_meta: dict,
                            gt_instances_ignore: Optional[InstanceData] = None,
                            unmap_outputs: bool = True) -> tuple:
        """Compute regression and classification targets for anchors in a
        single image.

        Args:
            flat_anchors (Tensor or :obj:`BaseBoxes`): Multi-level anchors
                of the image, which are concatenated into a single tensor
                or box type of shape (num_anchors, 4)
            valid_flags (Tensor): Multi level valid flags of the image,
                which are concatenated into a single tensor of
                    shape (num_anchors, ).
            gt_instances (:obj:`InstanceData`): Ground truth of instance
                annotations. It should includes ``bboxes`` and ``labels``
                attributes.
            img_meta (dict): Meta information for current image.
            gt_instances_ignore (:obj:`InstanceData`, optional): Instances
                to be ignored during training. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
                Defaults to None.
            unmap_outputs (bool): Whether to map outputs back to the original
                set of anchors.  Defaults to True.

        Returns:
            tuple:

                - labels (Tensor): Labels of each level.
                - label_weights (Tensor): Label weights of each level.
                - bbox_targets (Tensor): BBox targets of each level.
                - bbox_weights (Tensor): BBox weights of each level.
                - pos_inds (Tensor): positive samples indexes.
                - neg_inds (Tensor): negative samples indexes.
                - sampling_result (:obj:`SamplingResult`): Sampling results.
        """
        inside_flags = anchor_inside_flags(flat_anchors, valid_flags,
                                           img_meta['img_shape'][:2],
                                           self.train_cfg['allowed_border'])
        if not inside_flags.any():
            raise ValueError(
                'There is no valid anchor inside the image boundary. Please '
                'check the image size and anchor sizes, or set '
                '``allowed_border`` to -1 to skip the condition.')
        # assign gt and sample anchors
        anchors = flat_anchors[inside_flags]

        pred_instances = InstanceData(priors=anchors)
        assign_result = self.assigner.assign(pred_instances, gt_instances,
                                             gt_instances_ignore)
        # No sampling is required except for RPN and
        # Guided Anchoring algorithms
        sampling_result = self.sampler.sample(assign_result, pred_instances,
                                              gt_instances)

        num_valid_anchors = anchors.shape[0]
        bbox_targets = jt.zeros(num_valid_anchors, 4, dtype=anchors.dtype)
        bbox_weights = jt.zeros(num_valid_anchors, 4, dtype=anchors.dtype)

        # TODO: Considering saving memory, is it necessary to be long?
        labels = jt.full((num_valid_anchors, ),
                         self.num_classes,
                         dtype=jt.int64)
        label_weights = jt.zeros(num_valid_anchors, dtype=jt.float)

        pos_inds = sampling_result.pos_inds
        neg_inds = sampling_result.neg_inds
        # `bbox_coder.encode` accepts tensor or box type inputs and generates
        # tensor targets. If regressing decoded boxes, the code will convert
        # box type `pos_bbox_targets` to tensor.
        if len(pos_inds) > 0:
            if not self.reg_decoded_bbox:
                pos_bbox_targets = self.bbox_coder.encode(
                    sampling_result.pos_priors, sampling_result.pos_gt_bboxes)
            bbox_targets[pos_inds, :] = pos_bbox_targets
            bbox_weights[pos_inds, :] = 1.0

            labels[pos_inds] = sampling_result.pos_gt_labels
            if self.train_cfg['pos_weight'] <= 0:
                label_weights[pos_inds] = 1.0
            else:
                label_weights[pos_inds] = self.train_cfg['pos_weight']
        if len(neg_inds) > 0:
            label_weights[neg_inds] = 1.0

        # map up to original set of anchors
        if unmap_outputs:
            num_total_anchors = flat_anchors.size(0)
            labels = unmap(
                labels, num_total_anchors, inside_flags,
                fill=self.num_classes)  # fill bg label
            label_weights = unmap(label_weights, num_total_anchors,
                                  inside_flags)
            bbox_targets = unmap(bbox_targets, num_total_anchors, inside_flags)
            bbox_weights = unmap(bbox_weights, num_total_anchors, inside_flags)

        return (labels, label_weights, bbox_targets, bbox_weights, pos_inds,
                neg_inds, sampling_result)

    def get_targets(self,
                    anchor_list: List[List[jt.Var]],
                    valid_flag_list: List[List[jt.Var]],
                    batch_gt_instances: InstanceList,
                    batch_img_metas: List[dict],
                    batch_gt_instances_ignore: OptInstanceList = None,
                    unmap_outputs: bool = True,
                    return_sampling_results: bool = False) -> tuple:
        """Compute regression and classification targets for anchors in
        multiple images.

        Args:
            anchor_list (list[list[Tensor]]): Multi level anchors of each
                image. The outer list indicates images, and the inner list
                corresponds to feature levels of the image. Each element of
                the inner list is a tensor of shape (num_anchors, 4).
            valid_flag_list (list[list[Tensor]]): Multi level valid flags of
                each image. The outer list indicates images, and the inner list
                corresponds to feature levels of the image. Each element of
                the inner list is a tensor of shape (num_anchors, )
            batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance. It usually includes ``bboxes`` and ``labels``
                attributes.
            batch_img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            batch_gt_instances_ignore (list[:obj:`InstanceData`], optional):
                Batch of gt_instances_ignore. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
                Defaults to None.
            unmap_outputs (bool): Whether to map outputs back to the original
                set of anchors. Defaults to True.
            return_sampling_results (bool): Whether to return the sampling
                results. Defaults to False.

        Returns:
            tuple: Usually returns a tuple containing learning targets.

                - labels_list (list[Tensor]): Labels of each level.
                - label_weights_list (list[Tensor]): Label weights of each
                  level.
                - bbox_targets_list (list[Tensor]): BBox targets of each level.
                - bbox_weights_list (list[Tensor]): BBox weights of each level.
                - avg_factor (int): Average factor that is used to average
                  the loss. When using sampling method, avg_factor is usually
                  the sum of positive and negative priors. When using
                  `PseudoSampler`, `avg_factor` is usually equal to the number
                  of positive priors.

            additional_returns: This function enables user-defined returns from
                `self._get_targets_single`. These returns are currently refined
                to properties at each feature map (i.e. having HxW dimension).
                The results will be concatenated after the end
        """
        num_imgs = len(batch_img_metas)
        assert len(anchor_list) == len(valid_flag_list) == num_imgs

        if batch_gt_instances_ignore is None:
            batch_gt_instances_ignore = [None] * num_imgs

        # anchor number of multi levels
        num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]]
        # concat all level anchors to a single tensor
        concat_anchor_list = []
        concat_valid_flag_list = []
        for i in range(num_imgs):
            assert len(anchor_list[i]) == len(valid_flag_list[i])
            concat_anchor_list.append(jt.concat(anchor_list[i]))
            concat_valid_flag_list.append(jt.concat(valid_flag_list[i]))

        # compute targets for each image
        results = multi_apply(
            self._get_targets_single,
            concat_anchor_list,
            concat_valid_flag_list,
            batch_gt_instances,
            batch_img_metas,
            batch_gt_instances_ignore,
            unmap_outputs=unmap_outputs)
        (all_labels, all_label_weights, all_bbox_targets, all_bbox_weights,
         pos_inds_list, neg_inds_list, sampling_results_list) = results[:7]
        rest_results = list(results[7:])  # user-added return values
        # Get `avg_factor` of all images, which calculate in `SamplingResult`.
        # When using sampling method, avg_factor is usually the sum of
        # positive and negative priors. When using `PseudoSampler`,
        # `avg_factor` is usually equal to the number of positive priors.
        avg_factor = sum(
            [results.avg_factor for results in sampling_results_list])
        # split targets to a list w.r.t. multiple levels
        labels_list = images_to_levels(all_labels, num_level_anchors)
        label_weights_list = images_to_levels(all_label_weights,
                                              num_level_anchors)
        bbox_targets_list = images_to_levels(all_bbox_targets,
                                             num_level_anchors)
        bbox_weights_list = images_to_levels(all_bbox_weights,
                                             num_level_anchors)
        res = (labels_list, label_weights_list, bbox_targets_list,
               bbox_weights_list, avg_factor)
        if return_sampling_results:
            res = res + (sampling_results_list, )
        for i, r in enumerate(rest_results):  # user-added return values
            rest_results[i] = images_to_levels(r, num_level_anchors)

        return res + tuple(rest_results)

    def loss_by_feat_single(self, cls_score: jt.Var, bbox_pred: jt.Var,
                            anchors: jt.Var, labels: jt.Var,
                            label_weights: jt.Var, bbox_targets: jt.Var,
                            bbox_weights: jt.Var, avg_factor: int) -> tuple:
        """Calculate the loss of a single scale level based on the features
        extracted by the detection head.

        Args:
            cls_score (Tensor): Box scores for each scale level
                Has shape (N, num_anchors * num_classes, H, W).
            bbox_pred (Tensor): Box energies / deltas for each scale
                level with shape (N, num_anchors * 4, H, W).
            anchors (Tensor): Box reference for each scale level with shape
                (N, num_total_anchors, 4).
            labels (Tensor): Labels of each anchors with shape
                (N, num_total_anchors).
            label_weights (Tensor): Label weights of each anchor with shape
                (N, num_total_anchors)
            bbox_targets (Tensor): BBox regression targets of each anchor
                weight shape (N, num_total_anchors, 4).
            bbox_weights (Tensor): BBox regression loss weights of each anchor
                with shape (N, num_total_anchors, 4).
            avg_factor (int): Average factor that is used to average the loss.

        Returns:
            tuple: loss components.
        """
        # classification loss
        labels = labels.reshape(-1)
        label_weights = label_weights.reshape(-1)
        cls_score = cls_score.permute(0, 2, 3,
                                      1).reshape(-1, self.cls_out_channels)
        loss_cls = self.loss_cls(
            cls_score, labels, label_weights, avg_factor=avg_factor)
        # regression loss
        target_dim = bbox_targets.size(-1)
        bbox_targets = bbox_targets.reshape(-1, target_dim)
        bbox_weights = bbox_weights.reshape(-1, target_dim)
        bbox_pred = bbox_pred.permute(0, 2, 3, 1).reshape(-1, 4)
        if self.reg_decoded_bbox:
            # When the regression loss (e.g. `IouLoss`, `GIouLoss`)
            # is applied directly on the decoded bounding boxes, it
            # decodes the already encoded coordinates to absolute format.
            anchors = anchors.reshape(-1, anchors.size(-1))
            bbox_pred = self.bbox_coder.decode(anchors, bbox_pred)
        loss_bbox = self.loss_bbox(
            bbox_pred, bbox_targets, bbox_weights, avg_factor=avg_factor)
        return loss_cls, loss_bbox

    def loss_by_feat(
            self,
            cls_scores: List[jt.Var],
            bbox_preds: List[jt.Var],
            batch_gt_instances: InstanceList,
            batch_img_metas: List[dict],
            batch_gt_instances_ignore: OptInstanceList = None) -> dict:
        """Calculate the loss based on the features extracted by the detection
        head.

        Args:
            cls_scores (list[Tensor]): Box scores for each scale level
                has shape (N, num_anchors * num_classes, H, W).
            bbox_preds (list[Tensor]): Box energies / deltas for each scale
                level with shape (N, num_anchors * 4, H, W).
            batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance. It usually includes ``bboxes`` and ``labels``
                attributes.
            batch_img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            batch_gt_instances_ignore (list[:obj:`InstanceData`], optional):
                Batch of gt_instances_ignore. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
                Defaults to None.

        Returns:
            dict: A dictionary of loss components.
        """
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
        assert len(featmap_sizes) == self.prior_generator.num_levels

        anchor_list, valid_flag_list = self.get_priors(featmap_sizes,
                                                       batch_img_metas)
        cls_reg_targets = self.get_targets(
            anchor_list,
            valid_flag_list,
            batch_gt_instances,
            batch_img_metas,
            batch_gt_instances_ignore=batch_gt_instances_ignore)
        (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list,
         avg_factor) = cls_reg_targets

        # anchor number of multi levels
        num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]]
        # concat all level anchors and flags to a single tensor
        concat_anchor_list = []
        for i in range(len(anchor_list)):
            concat_anchor_list.append(jt.concat(anchor_list[i]))
        all_anchor_list = images_to_levels(concat_anchor_list,
                                           num_level_anchors)

        losses_cls, losses_bbox = multi_apply(
            self.loss_by_feat_single,
            cls_scores,
            bbox_preds,
            all_anchor_list,
            labels_list,
            label_weights_list,
            bbox_targets_list,
            bbox_weights_list,
            avg_factor=avg_factor)
        return dict(loss_cls=losses_cls, loss_bbox=losses_bbox)

jittordet/models/dense_heads/anchor_head.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/dense_heads/base_dense_head.py
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from abc import ABCMeta, abstractmethod
from typing import List, Optional, Tuple

import jittor as jt
import jittor.nn as nn

from jittordet.engine import ConfigDict
from jittordet.models.utils import (batched_nms, filter_scores_and_topk,
                                    select_single_mlvl, unpack_gt_instances)
from jittordet.structures import InstanceData, InstanceList, SampleList


class BaseDenseHead(nn.Module, metaclass=ABCMeta):
    """Base class for DenseHeads.

    1. The ``init_weights`` method is used to initialize densehead's
    model parameters. After detector initialization, ``init_weights``
    is triggered when ``detector.init_weights()`` is called externally.

    2. The ``loss`` method is used to calculate the loss of densehead,
    which includes two steps: (1) the densehead model performs forward
    propagation to obtain the feature maps (2) The ``loss_by_feat`` method
    is called based on the feature maps to calculate the loss.

    .. code:: text

    loss(): forward() -> loss_by_feat()

    3. The ``predict`` method is used to predict detection results,
    which includes two steps: (1) the densehead model performs forward
    propagation to obtain the feature maps (2) The ``predict_by_feat`` method
    is called based on the feature maps to predict detection results including
    post-processing.

    .. code:: text

    predict(): forward() -> predict_by_feat()

    4. The ``loss_and_predict`` method is used to return loss and detection
    results at the same time. It will call densehead's ``forward``,
    ``loss_by_feat`` and ``predict_by_feat`` methods in order.  If one-stage is
    used as RPN, the densehead needs to return both losses and predictions.
    This predictions is used as the proposal of roihead.

    .. code:: text

    loss_and_predict(): forward() -> loss_by_feat() -> predict_by_feat()
    """

    def __init__(self) -> None:
        super().__init__()

    def loss(self, x: Tuple[jt.Var], batch_data_samples: SampleList) -> dict:
        """Perform forward propagation and loss calculation of the detection
        head on the features of the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.
            batch_data_samples (List[:obj:`DetDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.

        Returns:
            dict: A dictionary of loss components.
        """
        outs = self(x)

        outputs = unpack_gt_instances(batch_data_samples)
        (batch_gt_instances, batch_img_metas,
         batch_gt_instances_ignore) = outputs

        loss_inputs = outs + (batch_gt_instances, batch_img_metas,
                              batch_gt_instances_ignore)
        losses = self.loss_by_feat(*loss_inputs)
        return losses

    @abstractmethod
    def loss_by_feat(self, **kwargs) -> dict:
        """Calculate the loss based on the features extracted by the detection
        head."""
        pass

    def loss_and_predict(
        self,
        x: Tuple[jt.Var],
        batch_data_samples: SampleList,
        proposal_cfg: Optional[ConfigDict] = None
    ) -> Tuple[dict, InstanceList]:
        """Perform forward propagation of the head, then calculate loss and
        predictions from the features and data samples.

        Args:
            x (tuple[Tensor]): Features from FPN.
            batch_data_samples (list[:obj:`DetDataSample`]): Each item contains
                the meta information of each image and corresponding
                annotations.
            proposal_cfg (ConfigDict, optional): Test / postprocessing
                configuration, if None, test_cfg would be used.
                Defaults to None.

        Returns:
            tuple: the return value is a tuple contains:

                - losses: (dict[str, Tensor]): A dictionary of loss components.
                - predictions (list[:obj:`InstanceData`]): Detection
                  results of each image after the post process.
        """
        outputs = unpack_gt_instances(batch_data_samples)
        (batch_gt_instances, batch_img_metas,
         batch_gt_instances_ignore) = outputs

        outs = self(x)

        loss_inputs = outs + (batch_gt_instances, batch_img_metas,
                              batch_gt_instances_ignore)
        losses = self.loss_by_feat(*loss_inputs)

        predictions = self.predict_by_feat(
            *outs, batch_img_metas=batch_img_metas, cfg=proposal_cfg)
        return losses, predictions

    def predict(self,
                x: Tuple[jt.Var],
                batch_data_samples: SampleList,
                rescale: bool = False) -> InstanceList:
        """Perform forward propagation of the detection head and predict
        detection results on the features of the upstream network.

        Args:
            x (tuple[Tensor]): Multi-level features from the
                upstream network, each is a 4D-tensor.
            batch_data_samples (List[:obj:`DetDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
            rescale (bool, optional): Whether to rescale the results.
                Defaults to False.

        Returns:
            list[obj:`InstanceData`]: Detection results of each image
            after the post process.
        """
        batch_img_metas = [
            data_samples.metainfo for data_samples in batch_data_samples
        ]

        outs = self(x)

        predictions = self.predict_by_feat(
            *outs, batch_img_metas=batch_img_metas, rescale=rescale)
        return predictions

    def predict_by_feat(self,
                        cls_scores: List[jt.Var],
                        bbox_preds: List[jt.Var],
                        score_factors: Optional[List[jt.Var]] = None,
                        batch_img_metas: Optional[List[dict]] = None,
                        cfg: Optional[ConfigDict] = None,
                        rescale: bool = False,
                        with_nms: bool = True) -> InstanceList:
        """Transform a batch of output features extracted from the head into
        bbox results.

        Note: When score_factors is not None, the cls_scores are
        usually multiplied by it then obtain the real score used in NMS,
        such as CenterNess in FCOS, IoU branch in ATSS.

        Args:
            cls_scores (list[Tensor]): Classification scores for all
                scale levels, each is a 4D-tensor, has shape
                (batch_size, num_priors * num_classes, H, W).
            bbox_preds (list[Tensor]): Box energies / deltas for all
                scale levels, each is a 4D-tensor, has shape
                (batch_size, num_priors * 4, H, W).
            score_factors (list[Tensor], optional): Score factor for
                all scale level, each is a 4D-tensor, has shape
                (batch_size, num_priors * 1, H, W). Defaults to None.
            batch_img_metas (list[dict], Optional): Batch image meta info.
                Defaults to None.
            cfg (ConfigDict, optional): Test / postprocessing
                configuration, if None, test_cfg would be used.
                Defaults to None.
            rescale (bool): If True, return boxes in original image space.
                Defaults to False.
            with_nms (bool): If True, do nms before return boxes.
                Defaults to True.

        Returns:
            list[:obj:`InstanceData`]: Object detection results of each image
            after the post process. Each item usually contains following keys.

                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        assert len(cls_scores) == len(bbox_preds)

        if score_factors is None:
            # e.g. Retina, FreeAnchor, Foveabox, etc.
            with_score_factors = False
        else:
            # e.g. FCOS, PAA, ATSS, AutoAssign, etc.
            with_score_factors = True
            assert len(cls_scores) == len(score_factors)

        num_levels = len(cls_scores)

        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
        mlvl_priors = self.prior_generator.grid_priors(
            featmap_sizes, dtype=cls_scores[0].dtype)

        result_list = []

        for img_id in range(len(batch_img_metas)):
            img_meta = batch_img_metas[img_id]
            cls_score_list = select_single_mlvl(
                cls_scores, img_id, detach=True)
            bbox_pred_list = select_single_mlvl(
                bbox_preds, img_id, detach=True)
            if with_score_factors:
                score_factor_list = select_single_mlvl(
                    score_factors, img_id, detach=True)
            else:
                score_factor_list = [None for _ in range(num_levels)]

            results = self._predict_by_feat_single(
                cls_score_list=cls_score_list,
                bbox_pred_list=bbox_pred_list,
                score_factor_list=score_factor_list,
                mlvl_priors=mlvl_priors,
                img_meta=img_meta,
                cfg=cfg,
                rescale=rescale,
                with_nms=with_nms)
            result_list.append(results)
        return result_list

    def _predict_by_feat_single(self,
                                cls_score_list: List[jt.Var],
                                bbox_pred_list: List[jt.Var],
                                score_factor_list: List[jt.Var],
                                mlvl_priors: List[jt.Var],
                                img_meta: dict,
                                cfg: ConfigDict,
                                rescale: bool = False,
                                with_nms: bool = True) -> InstanceData:
        """Transform a single image's features extracted from the head into
        bbox results.

        Args:
            cls_score_list (list[Tensor]): Box scores from all scale
                levels of a single image, each item has shape
                (num_priors * num_classes, H, W).
            bbox_pred_list (list[Tensor]): Box energies / deltas from
                all scale levels of a single image, each item has shape
                (num_priors * 4, H, W).
            score_factor_list (list[Tensor]): Score factor from all scale
                levels of a single image, each item has shape
                (num_priors * 1, H, W).
            mlvl_priors (list[Tensor]): Each element in the list is
                the priors of a single level in feature pyramid. In all
                anchor-based methods, it has shape (num_priors, 4). In
                all anchor-free methods, it has shape (num_priors, 2)
                when `with_stride=True`, otherwise it still has shape
                (num_priors, 4).
            img_meta (dict): Image meta info.
            cfg (mmengine.Config): Test / postprocessing configuration,
                if None, test_cfg would be used.
            rescale (bool): If True, return boxes in original image space.
                Defaults to False.
            with_nms (bool): If True, do nms before return boxes.
                Defaults to True.

        Returns:
            :obj:`InstanceData`: Detection results of each image
            after the post process.
            Each item usually contains following keys.

                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        if score_factor_list[0] is None:
            # e.g. Retina, FreeAnchor, etc.
            with_score_factors = False
        else:
            # e.g. FCOS, PAA, ATSS, etc.
            with_score_factors = True

        cfg = self.test_cfg if cfg is None else cfg
        cfg = copy.deepcopy(cfg)
        img_shape = img_meta['img_shape']
        nms_pre = cfg.get('nms_pre', -1)

        mlvl_bbox_preds = []
        mlvl_valid_priors = []
        mlvl_scores = []
        mlvl_labels = []
        if with_score_factors:
            mlvl_score_factors = []
        else:
            mlvl_score_factors = None
        for level_idx, (cls_score, bbox_pred, score_factor, priors) in \
                enumerate(zip(cls_score_list, bbox_pred_list,
                              score_factor_list, mlvl_priors)):

            assert cls_score.size()[-2:] == bbox_pred.size()[-2:]

            bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, 4)
            if with_score_factors:
                score_factor = score_factor.permute(1, 2,
                                                    0).reshape(-1).sigmoid()
            cls_score = cls_score.permute(1, 2,
                                          0).reshape(-1, self.cls_out_channels)
            if self.use_sigmoid_cls:
                scores = cls_score.sigmoid()
            else:
                # remind that we set FG labels to [0, num_class-1]
                # since mmdet v2.0
                # BG cat_id: num_class
                scores = cls_score.softmax(-1)[:, :-1]

            # After https://github.com/open-mmlab/mmdetection/pull/6268/,
            # this operation keeps fewer bboxes under the same `nms_pre`.
            # There is no difference in performance for most models. If you
            # find a slight drop in performance, you can set a larger
            # `nms_pre` than before.
            score_thr = cfg.get('score_thr', 0)

            results = filter_scores_and_topk(
                scores, score_thr, nms_pre,
                dict(bbox_pred=bbox_pred, priors=priors))
            scores, labels, keep_idxs, filtered_results = results

            bbox_pred = filtered_results['bbox_pred']
            priors = filtered_results['priors']

            if with_score_factors:
                score_factor = score_factor[keep_idxs]

            mlvl_bbox_preds.append(bbox_pred)
            mlvl_valid_priors.append(priors)
            mlvl_scores.append(scores)
            mlvl_labels.append(labels)

            if with_score_factors:
                mlvl_score_factors.append(score_factor)

        bbox_pred = jt.concat(mlvl_bbox_preds)
        priors = jt.concat(mlvl_valid_priors)
        bboxes = self.bbox_coder.decode(priors, bbox_pred, max_shape=img_shape)

        results = InstanceData()
        results.bboxes = bboxes
        results.scores = jt.concat(mlvl_scores)
        results.labels = jt.concat(mlvl_labels)
        if with_score_factors:
            results.score_factors = jt.concat(mlvl_score_factors)

        return self._bbox_post_process(
            results=results,
            cfg=cfg,
            rescale=rescale,
            with_nms=with_nms,
            img_meta=img_meta)

    def _bbox_post_process(self,
                           results: InstanceData,
                           cfg: ConfigDict,
                           rescale: bool = False,
                           with_nms: bool = True,
                           img_meta: Optional[dict] = None) -> InstanceData:
        """bbox post-processing method.

        The boxes would be rescaled to the original image scale and do
        the nms operation. Usually `with_nms` is False is used for aug test.

        Args:
            results (:obj:`InstaceData`): Detection instance results,
                each item has shape (num_bboxes, ).
            cfg (ConfigDict): Test / postprocessing configuration,
                if None, test_cfg would be used.
            rescale (bool): If True, return boxes in original image space.
                Default to False.
            with_nms (bool): If True, do nms before return boxes.
                Default to True.
            img_meta (dict, optional): Image meta info. Defaults to None.

        Returns:
            :obj:`InstanceData`: Detection results of each image
            after the post process.
            Each item usually contains following keys.

                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        if rescale:
            assert img_meta.get('scale_factor') is not None
            results.bboxes /= jt.array(
                img_meta['scale_factor'], dtype=results.bboxes.dtype).repeat(
                    (1, 2))

        if hasattr(results, 'score_factors'):
            # TODO： Add sqrt operation in order to be consistent with
            #  the paper.
            score_factors = results.pop('score_factors')
            results.scores = results.scores * score_factors

        # filter small size bboxes
        if cfg.get('min_bbox_size', -1) >= 0:
            w = results.bboxes[:, 2] - results.bboxes[:, 0]
            h = results.bboxes[:, 3] - results.bboxes[:, 1]
            valid_mask = (w > cfg.min_bbox_size) & (h > cfg.min_bbox_size)
            if not valid_mask.all():
                results = results[valid_mask]

        # TODO: deal with `with_nms` and `nms_cfg=None` in test_cfg
        if with_nms and results.bboxes.numel() > 0:
            _, det_scores, keep_idxs = batched_nms(results.bboxes,
                                                   results.scores,
                                                   results.labels, cfg.nms)
            results = results[keep_idxs]
            # some nms would reweight the score, such as softnms
            results.scores = det_scores
            results = results[:cfg.max_per_img]

        return results

jittordet/models/dense_heads/base_dense_head.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/dense_heads/retina_head.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor.nn as nn

from jittordet.engine import MODELS
from ..layers import ConvModule
from ..utils import bias_init_with_prob, normal_init
from .anchor_head import AnchorHead


@MODELS.register_module()
class RetinaHead(AnchorHead):
    r"""An anchor-based head used in `RetinaNet
    <https://arxiv.org/pdf/1708.02002.pdf>`_.

    The head contains two subnetworks. The first classifies anchor boxes and
    the second regresses deltas for the anchors.

    Example:
        >>> import torch
        >>> self = RetinaHead(11, 7)
        >>> x = torch.rand(1, 7, 32, 32)
        >>> cls_score, bbox_pred = self.forward_single(x)
        >>> # Each anchor predicts a score for each class except background
        >>> cls_per_anchor = cls_score.shape[1] / self.num_anchors
        >>> box_per_anchor = bbox_pred.shape[1] / self.num_anchors
        >>> assert cls_per_anchor == (self.num_classes)
        >>> assert box_per_anchor == 4
    """

    def __init__(self,
                 num_classes,
                 in_channels,
                 stacked_convs=4,
                 conv_cfg=None,
                 norm_cfg=None,
                 anchor_generator=dict(
                     type='AnchorGenerator',
                     octave_base_scale=4,
                     scales_per_octave=3,
                     ratios=[0.5, 1.0, 2.0],
                     strides=[8, 16, 32, 64, 128]),
                 **kwargs):
        assert stacked_convs >= 0, \
            '`stacked_convs` must be non-negative integers, ' \
            f'but got {stacked_convs} instead.'
        self.stacked_convs = stacked_convs
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        super(RetinaHead, self).__init__(
            num_classes,
            in_channels,
            anchor_generator=anchor_generator,
            **kwargs)

    def _init_layers(self):
        """Initialize layers of the head."""
        self.relu = nn.ReLU(inplace=True)
        self.cls_convs = nn.ModuleList()
        self.reg_convs = nn.ModuleList()
        in_channels = self.in_channels
        for i in range(self.stacked_convs):
            self.cls_convs.append(
                ConvModule(
                    in_channels,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            self.reg_convs.append(
                ConvModule(
                    in_channels,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            in_channels = self.feat_channels
        self.retina_cls = nn.Conv2d(
            in_channels,
            self.num_base_priors * self.cls_out_channels,
            3,
            padding=1)
        self.retina_reg = nn.Conv2d(
            in_channels, self.num_base_priors * 4, 3, padding=1)

    def init_weights(self):
        for m in self.cls_convs.modules():
            if isinstance(m, nn.Conv):
                normal_init(m, std=0.01)
        for m in self.reg_convs.modules():
            if isinstance(m, nn.Conv):
                normal_init(m, std=0.01)

        normal_init(self.retina_reg, std=0.01)
        bias = bias_init_with_prob(0.01)
        normal_init(self.retina_cls, std=0.01, bias=bias)

    def execute_single(self, x):
        """Forward feature of a single scale level.

        Args:
            x (Tensor): Features of a single scale level.

        Returns:
            tuple:
                cls_score (Tensor): Cls scores for a single scale level
                    the channels number is num_anchors * num_classes.
                bbox_pred (Tensor): Box energies / deltas for a single scale
                    level, the channels number is num_anchors * 4.
        """
        cls_feat = x
        reg_feat = x
        for cls_conv in self.cls_convs:
            cls_feat = cls_conv(cls_feat)
        for reg_conv in self.reg_convs:
            reg_feat = reg_conv(reg_feat)
        cls_score = self.retina_cls(cls_feat)
        bbox_pred = self.retina_reg(reg_feat)
        return cls_score, bbox_pred

jittordet/models/dense_heads/retina_head.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/dense_heads/rpn_head.py
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from typing import List, Optional

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS, ConfigDict
from jittordet.structures import InstanceData, InstanceList, OptInstanceList
from ..layers import ConvModule
from ..utils import batched_nms, normal_init
from .anchor_head import AnchorHead


@MODELS.register_module()
class RPNHead(AnchorHead):

    def __init__(self, num_classes, in_channels, num_convs=1, **kwargs):
        self.num_convs = num_convs
        assert num_classes == 1
        super(RPNHead, self).__init__(num_classes, in_channels, **kwargs)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv):
                normal_init(m, std=0.01)

    def _init_layers(self):
        """Initialize layers of the head."""
        if self.num_convs > 1:
            rpn_convs = []
            for i in range(self.num_convs):
                if i == 0:
                    in_channels = self.in_channels
                else:
                    in_channels = self.feat_channels
                # use ``inplace=False`` to avoid error: one of the variables
                # needed for gradient computation has been modified by an
                # inplace operation.
                rpn_convs.append(
                    ConvModule(in_channels, self.feat_channels, 3, padding=1))
            self.rpn_conv = nn.Sequential(*rpn_convs)
        else:
            self.rpn_conv = nn.Conv2d(
                self.in_channels, self.feat_channels, 3, padding=1)
        self.rpn_cls = nn.Conv2d(self.feat_channels,
                                 self.num_base_priors * self.cls_out_channels,
                                 1)
        self.rpn_reg = nn.Conv2d(self.feat_channels, self.num_base_priors * 4,
                                 1)

    def execute_single(self, x):
        """Forward feature of a single scale level.

        Args:
            x (Tensor): Features of a single scale level.

        Returns:
            tuple:
                cls_score (Tensor): Cls scores for a single scale level
                    the channels number is num_anchors * num_classes.
                bbox_pred (Tensor): Box energies / deltas for a single scale
                    level, the channels number is num_anchors * 4.
        """
        x = self.rpn_conv(x)
        x = nn.relu(x)
        rpn_cls_score = self.rpn_cls(x)
        rpn_bbox_pred = self.rpn_reg(x)
        return rpn_cls_score, rpn_bbox_pred

    def loss_by_feat(self,
                     cls_scores: List[jt.Var],
                     bbox_preds: List[jt.Var],
                     batch_gt_instances: InstanceList,
                     batch_img_metas: List[dict],
                     batch_gt_instances_ignore: OptInstanceList = None) \
            -> dict:
        """Calculate the loss based on the features extracted by the detection
        head.
        Args:
            cls_scores (list[Tensor]): Box scores for each scale level,
                has shape (N, num_anchors * num_classes, H, W).
            bbox_preds (list[Tensor]): Box energies / deltas for each scale
                level with shape (N, num_anchors * 4, H, W).
            batch_gt_instances (list[obj:InstanceData]): Batch of gt_instance.
                It usually includes ``bboxes`` and ``labels`` attributes.
            batch_img_metas (list[dict]): Meta information of each image, e.g.,
                image size, scaling factor, etc.
            batch_gt_instances_ignore (list[obj:InstanceData], Optional):
                Batch of gt_instances_ignore. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        losses = super().loss_by_feat(
            cls_scores,
            bbox_preds,
            batch_gt_instances,
            batch_img_metas,
            batch_gt_instances_ignore=batch_gt_instances_ignore)
        return dict(
            loss_rpn_cls=losses['loss_cls'], loss_rpn_bbox=losses['loss_bbox'])

    def _predict_by_feat_single(self,
                                cls_score_list: List[jt.Var],
                                bbox_pred_list: List[jt.Var],
                                score_factor_list: List[jt.Var],
                                mlvl_priors: List[jt.Var],
                                img_meta: dict,
                                cfg: ConfigDict,
                                rescale: bool = False,
                                with_nms: bool = True) -> InstanceData:
        cfg = self.test_cfg if cfg is None else cfg
        cfg = copy.deepcopy(cfg)
        img_shape = img_meta['img_shape']
        nms_pre = cfg.get('nms_pre', -1)

        mlvl_bbox_preds = []
        mlvl_valid_priors = []
        mlvl_scores = []
        level_ids = []

        for level_idx, (cls_score, bbox_pred, priors) in enumerate(
                zip(cls_score_list, bbox_pred_list, mlvl_priors)):
            assert cls_score.shape[-2:] == bbox_pred.shape[-2:]
            bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, 4)
            cls_score = cls_score.permute(1, 2,
                                          0).reshape(-1, self.cls_out_channels)
            if self.use_sigmoid_cls:
                scores = cls_score.sigmoid()
            else:
                # remind that we set FG labels to [0] since mmdet v2.0
                # BG cat_id: 1
                scores = nn.softmax(cls_score, -1)[:, :-1]
            scores = jt.squeeze(scores, -1)
            if 0 < nms_pre < scores.shape[0]:
                # sort is faster than topk
                # _, topk_inds = scores.topk(cfg.nms_pre)
                rank_inds, ranked_scores = scores.argsort(descending=True)
                topk_inds = rank_inds[:nms_pre]
                scores = ranked_scores[:nms_pre]
                bbox_pred = bbox_pred[topk_inds, :]
                priors = priors[topk_inds]

            mlvl_bbox_preds.append(bbox_pred)
            mlvl_valid_priors.append(priors)
            mlvl_scores.append(scores)
            level_ids.append(
                jt.full((scores.size(0), ), level_idx, dtype=jt.int64))

        bbox_pred = jt.concat(mlvl_bbox_preds)
        priors = jt.concat(mlvl_valid_priors)
        bboxes = self.bbox_coder.decode(priors, bbox_pred, max_shape=img_shape)

        results = InstanceData()
        results.bboxes = bboxes
        results.scores = jt.concat(mlvl_scores)
        results.level_ids = jt.concat(level_ids)
        return self._bbox_post_process(
            results=results, cfg=cfg, rescale=rescale, img_meta=img_meta)

    def _bbox_post_process(self,
                           results: InstanceData,
                           cfg: ConfigDict,
                           rescale: bool = False,
                           with_nms: bool = True,
                           img_meta: Optional[dict] = None) -> InstanceData:
        assert with_nms, '`with_nms` must be True in RPNHead'
        if rescale:
            assert img_meta.get('scale_factor') is not None
            results.bboxes /= jt.array(
                img_meta['scale_factor'], dtype=results.bboxes.dtype).repeat(
                    (1, 2))

        # filter small size bboxes
        if cfg.get('min_bbox_size', -1) >= 0:
            w = results.bboxes[:, 2] - results.bboxes[:, 0]
            h = results.bboxes[:, 3] - results.bboxes[:, 1]
            valid_mask = (w > cfg.min_bbox_size) & (h > cfg.min_bbox_size)
            if not valid_mask.all():
                results = results[valid_mask]

        if results.bboxes.numel() > 0:
            bboxes = results.bboxes
            # TODO batched_nms
            _, det_scores, keep_idxs = batched_nms(bboxes, results.scores,
                                                   results.level_ids, cfg.nms)
            results = results[keep_idxs]
            # some nms would reweight the score, such as softnms
            results.scores = det_scores
            results = results[:cfg.max_per_img]
            # TODO: This would unreasonably show the 0th class label
            #  in visualization
            results.labels = jt.zeros(len(results), dtype=jt.int64)
            del results.level_ids
        else:
            # To avoid some potential error
            results_ = InstanceData()
            results_.bboxes = jt.zeros_like(results.bboxes)
            results_.scores = jt.zeros_like(results.scores)
            results_.labels = jt.zeros_like(results.scores)
            results = results_
        return results

jittordet/models/dense_heads/rpn_head.py
JittorCV/jittordet
from .base_framework import BaseFramework
from .multi_stage import MultiStageFramework
from .rpn import RPNFramework
from .single_stage import SingleStageFramework

__all__ = [
    'BaseFramework', 'SingleStageFramework', 'MultiStageFramework',
    'RPNFramework'
]

jittordet/models/frameworks/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/base.py
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod
from collections import OrderedDict
from typing import Dict, List, Tuple, Union

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS
from jittordet.structures import (DetDataSample, InstanceList, OptSampleList,
                                  SampleList)
from jittordet.utils import is_list_of

ForwardResults = Union[Dict[str, jt.Var], List[DetDataSample], Tuple[jt.Var],
                       jt.Var]


class BaseFramework(nn.Module, metaclass=ABCMeta):
    """Base class for detectors.

    Args:
       data_preprocessor (dict or ConfigDict, optional): The pre-process
           config of :class:`BaseDataPreprocessor`.  it usually includes,
            ``pad_size_divisor``, ``pad_value``, ``mean`` and ``std``.
       init_cfg (dict or ConfigDict, optional): the config to control the
           initialization. Defaults to None.
    """

    def __init__(self, preprocessor=None):
        super().__init__()
        self.preprocessor = MODELS.build(preprocessor)

    @property
    def with_neck(self) -> bool:
        """bool: whether the detector has a neck"""
        return hasattr(self, 'neck') and self.neck is not None

    # TODO: these properties need to be carefully handled
    # for both single stage & two stage detectors
    @property
    def with_shared_head(self) -> bool:
        """bool: whether the detector has a shared head in the RoI Head"""
        return hasattr(self, 'roi_head') and self.roi_head.with_shared_head

    @property
    def with_bbox(self) -> bool:
        """bool: whether the detector has a bbox head"""
        return ((hasattr(self, 'roi_head') and self.roi_head.with_bbox)
                or (hasattr(self, 'bbox_head') and self.bbox_head is not None))

    @property
    def with_mask(self) -> bool:
        """bool: whether the detector has a mask head"""
        return ((hasattr(self, 'roi_head') and self.roi_head.with_mask)
                or (hasattr(self, 'mask_head') and self.mask_head is not None))

    def execute(self, data, phase: str = 'tensor') -> ForwardResults:
        """The unified entry for a forward process in both training and test.

        The method should accept three modes: "tensor", "predict" and "loss":

        - "tensor": Forward the whole network and return tensor or tuple of
        tensor without any post-processing, same as a common nn.Module.
        - "predict": Forward and return the predictions, which are fully
        processed to a list of :obj:`DetDataSample`.
        - "loss": Forward and return a dict of losses according to the given
        inputs and data samples.

        Note that this method doesn't handle either back propagation or
        parameter update, which are supposed to be done in :meth:`train_step`.

        Args:
            inputs (torch.Tensor): The input tensor with shape
                (N, C, ...) in general.
            data_samples (list[:obj:`DetDataSample`], optional): A batch of
                data samples that contain annotations and predictions.
                Defaults to None.
            phase (str): Return what kind of value. Defaults to 'tensor'.

        Returns:
            The return type depends on ``mode``.

            - If ``mode="tensor"``, return a tensor or a tuple of tensor.
            - If ``mode="predict"``, return a list of :obj:`DetDataSample`.
            - If ``mode="loss"``, return a dict of tensor.
        """
        data = self.preprocessor(data)
        inputs, data_samples = data['inputs'], data['data_samples']
        if phase == 'loss':
            results = self.loss(inputs, data_samples)
            return self.parse_losses(results)
        elif phase == 'predict':
            return self.predict(inputs, data_samples)
        elif phase == 'tensor':
            return self._execute(inputs, data_samples)
        else:
            raise RuntimeError(f'Invalid mode "{phase}". '
                               'Only supports loss, predict and tensor mode')

    def parse_losses(
            self, losses: Dict[str,
                               jt.Var]) -> Tuple[jt.Var, Dict[str, jt.Var]]:
        """Parses the raw outputs (losses) of the network.

        Args:
            losses (dict): Raw output of the network, which usually contain
                losses and other necessary information.

        Returns:
            tuple[Tensor, dict]: There are two elements. The first is the
            loss tensor passed to optim_wrapper which may be a weighted sum of
            all losses, and the second is log_vars which will be sent to the
            logger.
        """
        log_vars = []
        for loss_name, loss_value in losses.items():
            if isinstance(loss_value, jt.Var):
                log_vars.append([loss_name, loss_value.mean()])
            elif is_list_of(loss_value, jt.Var):
                log_vars.append(
                    [loss_name,
                     sum(_loss.mean() for _loss in loss_value)])
            else:
                raise TypeError(
                    f'{loss_name} is not a tensor or list of tensors')

        loss = sum(value for key, value in log_vars if 'loss' in key)
        log_vars.insert(0, ['loss', loss])
        log_vars = OrderedDict(log_vars)  # type: ignore

        return loss, log_vars  # type: ignore

    @abstractmethod
    def loss(self, batch_inputs: jt.Var,
             batch_data_samples: List[DetDataSample]) -> Union[dict, tuple]:
        """Calculate losses from a batch of inputs and data samples."""
        pass

    @abstractmethod
    def predict(
            self, batch_inputs: jt.Var,
            batch_data_samples: List[DetDataSample]) -> List[DetDataSample]:
        """Predict results from a batch of inputs and data samples with post-
        processing."""
        pass

    @abstractmethod
    def _execute(self,
                 batch_inputs: jt.Var,
                 batch_data_samples: OptSampleList = None):
        """Network forward process.

        Usually includes backbone, neck and head forward without any post-
        processing.
        """
        pass

    @abstractmethod
    def extract_feat(self, batch_inputs: jt.Var):
        """Extract features from images."""
        pass

    def add_pred_to_datasample(self, data_samples: SampleList,
                               results_list: InstanceList) -> SampleList:
        """Add predictions to `DetDataSample`.

        Args:
            data_samples (list[:obj:`DetDataSample`], optional): A batch of
                data samples that contain annotations and predictions.
            results_list (list[:obj:`InstanceData`]): Detection results of
                each image.

        Returns:
            list[:obj:`DetDataSample`]: Detection results of the
            input images. Each DetDataSample usually contain
            'pred_instances'. And the ``pred_instances`` usually
            contains following keys.

                - scores (Tensor): Classification scores, has a shape
                    (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                    (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                    the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        for data_sample, pred_instances in zip(data_samples, results_list):
            data_sample.pred_instances = pred_instances
        return data_samples

jittordet/models/frameworks/base_framework.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/detectors/two_stage.py
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import warnings
from typing import Tuple

import jittor as jt

from jittordet.engine import MODELS, ConfigType, OptConfigType
from jittordet.structures import SampleList
from .base_framework import BaseFramework


@MODELS.register_module()
class MultiStageFramework(BaseFramework):
    """Base class for two-stage detectors.

    Two-stage detectors typically consisting of a region proposal network and a
    task-specific regression head.
    """

    def __init__(self,
                 backbone: ConfigType,
                 neck: OptConfigType = None,
                 rpn_head: OptConfigType = None,
                 roi_head: OptConfigType = None,
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None,
                 preprocessor: OptConfigType = None) -> None:
        super().__init__(preprocessor=preprocessor)
        self.backbone = MODELS.build(backbone)

        if neck is not None:
            self.neck = MODELS.build(neck)

        if rpn_head is not None:
            rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
            rpn_head_ = rpn_head.copy()
            rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn)
            rpn_head_num_classes = rpn_head_.get('num_classes', None)
            if rpn_head_num_classes is None:
                rpn_head_.update(num_classes=1)
            else:
                if rpn_head_num_classes != 1:
                    warnings.warn(
                        'The `num_classes` should be 1 in RPN, but get '
                        f'{rpn_head_num_classes}, please set '
                        'rpn_head.num_classes = 1 in your config file.')
                    rpn_head_.update(num_classes=1)
            self.rpn_head = MODELS.build(rpn_head_)

        if roi_head is not None:
            # update train and test cfg here for now
            # TODO: refactor assigner & sampler
            rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None
            roi_head.update(train_cfg=rcnn_train_cfg)
            roi_head.update(test_cfg=test_cfg.rcnn)
            self.roi_head = MODELS.build(roi_head)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    @property
    def with_rpn(self) -> bool:
        """bool: whether the detector has RPN"""
        return hasattr(self, 'rpn_head') and self.rpn_head is not None

    @property
    def with_roi_head(self) -> bool:
        """bool: whether the detector has a RoI head"""
        return hasattr(self, 'roi_head') and self.roi_head is not None

    def extract_feat(self, batch_inputs: jt.Var) -> Tuple[jt.Var]:
        """Extract features.

        Args:
            batch_inputs (Tensor): Image tensor with shape (N, C, H ,W).

        Returns:
            tuple[Tensor]: Multi-level features that may have
            different resolutions.
        """
        x = self.backbone(batch_inputs)
        if self.with_neck:
            x = self.neck(x)
        return x

    def _execute(self, batch_inputs: jt.Var,
                 batch_data_samples: SampleList) -> tuple:
        """Network forward process. Usually includes backbone, neck and head
        forward without any post-processing.

        Args:
            batch_inputs (Tensor): Inputs with shape (N, C, H, W).

        Returns:
            tuple: A tuple of features from ``rpn_head`` and ``roi_head``
            forward.
        """
        results = ()
        x = self.extract_feat(batch_inputs)

        if self.with_rpn:
            rpn_results_list = self.rpn_head.predict(
                x, batch_data_samples, rescale=False)
        else:
            assert batch_data_samples[0].get('proposals', None) is not None
            rpn_results_list = [
                data_sample.proposals for data_sample in batch_data_samples
            ]

        roi_outs = self.roi_head.forward(x, rpn_results_list)
        results = results + (roi_outs, )
        return results

    def loss(self, batch_inputs: jt.Var,
             batch_data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            batch_inputs (Tensor): Input images of shape (N, C, H, W).
                These should usually be mean centered and std scaled.
            batch_data_samples (List[:obj:`DetDataSample`]): The batch
                data samples. It usually includes information such
                as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`.

        Returns:
            dict: A dictionary of loss components
        """
        x = self.extract_feat(batch_inputs)

        losses = dict()

        # RPN forward and loss
        if self.with_rpn:
            proposal_cfg = self.train_cfg.get('rpn_proposal',
                                              self.test_cfg.rpn)
            rpn_data_samples = copy.deepcopy(batch_data_samples)
            # set cat_id of gt_labels to 0 in RPN
            for data_sample in rpn_data_samples:
                data_sample.gt_instances.labels = \
                    jt.zeros_like(data_sample.gt_instances.labels)

            rpn_losses, rpn_results_list = self.rpn_head.loss_and_predict(
                x, rpn_data_samples, proposal_cfg=proposal_cfg)
            # avoid get same name with roi_head loss
            keys = rpn_losses.keys()
            for key in list(keys):
                if 'loss' in key and 'rpn' not in key:
                    rpn_losses[f'rpn_{key}'] = rpn_losses.pop(key)
            losses.update(rpn_losses)
        else:
            assert batch_data_samples[0].get('proposals', None) is not None
            # use pre-defined proposals in InstanceData for the second stage
            # to extract ROI features.
            rpn_results_list = [
                data_sample.proposals for data_sample in batch_data_samples
            ]

        roi_losses = self.roi_head.loss(x, rpn_results_list,
                                        batch_data_samples)
        losses.update(roi_losses)

        return losses

    def predict(self,
                batch_inputs: jt.Var,
                batch_data_samples: SampleList,
                rescale: bool = True) -> SampleList:
        """Predict results from a batch of inputs and data samples with post-
        processing.

        Args:
            batch_inputs (Tensor): Inputs with shape (N, C, H, W).
            batch_data_samples (List[:obj:`DetDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
            rescale (bool): Whether to rescale the results.
                Defaults to True.

        Returns:
            list[:obj:`DetDataSample`]: Return the detection results of the
            input images. The returns value is DetDataSample,
            which usually contain 'pred_instances'. And the
            ``pred_instances`` usually contains following keys.

                - scores (Tensor): Classification scores, has a shape
                    (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                    (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                    the last dimension 4 arrange as (x1, y1, x2, y2).
                - masks (Tensor): Has a shape (num_instances, H, W).
        """

        assert self.with_bbox, 'Bbox head must be implemented.'
        x = self.extract_feat(batch_inputs)

        # If there are no pre-defined proposals, use RPN to get proposals
        if batch_data_samples[0].get('proposals', None) is None:
            rpn_results_list = self.rpn_head.predict(
                x, batch_data_samples, rescale=False)
        else:
            rpn_results_list = [
                data_sample.proposals for data_sample in batch_data_samples
            ]

        results_list = self.roi_head.predict(
            x, rpn_results_list, batch_data_samples, rescale=rescale)

        batch_data_samples = self.add_pred_to_datasample(
            batch_data_samples, results_list)
        return batch_data_samples

jittordet/models/frameworks/multi_stage.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import warnings

import jittor as jt

from jittordet.engine import MODELS, ConfigType, OptConfigType
from jittordet.structures import SampleList
from .single_stage import SingleStageFramework


@MODELS.register_module()
class RPNFramework(SingleStageFramework):
    """Implementation of Region Proposal Network.

    Args:
        backbone (:obj:`ConfigDict` or dict): The backbone config.
        neck (:obj:`ConfigDict` or dict): The neck config.
        bbox_head (:obj:`ConfigDict` or dict): The bbox head config.
        train_cfg (:obj:`ConfigDict` or dict, optional): The training config.
        test_cfg (:obj:`ConfigDict` or dict, optional): The testing config.
        data_preprocessor (:obj:`ConfigDict` or dict, optional): Config of
            :class:`DetDataPreprocessor` to process the input data.
            Defaults to None.
        init_cfg (:obj:`ConfigDict` or list[:obj:`ConfigDict`] or dict or
            list[dict], optional): Initialization config dict.
            Defaults to None.
    """

    def __init__(self,
                 backbone: ConfigType,
                 neck: ConfigType,
                 rpn_head: ConfigType,
                 train_cfg: ConfigType,
                 test_cfg: ConfigType,
                 preprocessor: OptConfigType = None,
                 **kwargs) -> None:
        super(SingleStageFramework, self).__init__(preprocessor=preprocessor)
        self.backbone = MODELS.build(backbone)
        self.neck = MODELS.build(neck) if neck is not None else None
        rpn_train_cfg = train_cfg['rpn'] if train_cfg is not None else None
        rpn_head_num_classes = rpn_head.get('num_classes', 1)
        if rpn_head_num_classes != 1:
            warnings.warn('The `num_classes` should be 1 in RPN, but get '
                          f'{rpn_head_num_classes}, please set '
                          'rpn_head.num_classes = 1 in your config file.')
            rpn_head.update(num_classes=1)
        rpn_head.update(train_cfg=rpn_train_cfg)
        rpn_head.update(test_cfg=test_cfg['rpn'])
        self.bbox_head = MODELS.build(rpn_head)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    def loss(self, batch_inputs: jt.Var,
             batch_data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            batch_inputs (Tensor): Input images of shape (N, C, H, W).
                These should usually be mean centered and std scaled.
            batch_data_samples (list[:obj:`DetDataSample`]): The batch
                data samples. It usually includes information such
                as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`.

        Returns:
            dict[str, Tensor]: A dictionary of loss components.
        """
        x = self.extract_feat(batch_inputs)

        # set cat_id of gt_labels to 0 in RPN
        rpn_data_samples = copy.deepcopy(batch_data_samples)
        for data_sample in rpn_data_samples:
            data_sample.gt_instances.labels = \
                jt.zeros_like(data_sample.gt_instances.labels)

        losses = self.bbox_head.loss(x, rpn_data_samples)
        return losses

jittordet/models/frameworks/rpn.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/detectors/single_stage.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple, Union

import jittor as jt

from jittordet.engine import MODELS, ConfigType, OptConfigType
from jittordet.structures import OptSampleList, SampleList
from .base_framework import BaseFramework


@MODELS.register_module()
class SingleStageFramework(BaseFramework):
    """Base class for single-stage detectors.

    Single-stage detectors directly and densely predict bounding boxes on the
    output features of the backbone+neck.
    """

    def __init__(self,
                 backbone: ConfigType,
                 neck: OptConfigType = None,
                 bbox_head: OptConfigType = None,
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None,
                 preprocessor: OptConfigType = None) -> None:
        super().__init__(preprocessor=preprocessor)
        self.backbone = MODELS.build(backbone)
        if neck is not None:
            self.neck = MODELS.build(neck)
        bbox_head.update(train_cfg=train_cfg)
        bbox_head.update(test_cfg=test_cfg)
        self.bbox_head = MODELS.build(bbox_head)
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

    def loss(self, batch_inputs: jt.Var,
             batch_data_samples: SampleList) -> Union[dict, list]:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            batch_inputs (Tensor): Input images of shape (N, C, H, W).
                These should usually be mean centered and std scaled.
            batch_data_samples (list[:obj:`DetDataSample`]): The batch
                data samples. It usually includes information such
                as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`.

        Returns:
            dict: A dictionary of loss components.
        """
        x = self.extract_feat(batch_inputs)
        losses = self.bbox_head.loss(x, batch_data_samples)
        return losses

    def predict(self,
                batch_inputs: jt.Var,
                batch_data_samples: SampleList,
                rescale: bool = True) -> SampleList:
        """Predict results from a batch of inputs and data samples with post-
        processing.

        Args:
            batch_inputs (Tensor): Inputs with shape (N, C, H, W).
            batch_data_samples (List[:obj:`DetDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
            rescale (bool): Whether to rescale the results.
                Defaults to True.

        Returns:
            list[:obj:`DetDataSample`]: Detection results of the
            input images. Each DetDataSample usually contain
            'pred_instances'. And the ``pred_instances`` usually
            contains following keys.

                - scores (Tensor): Classification scores, has a shape
                    (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                    (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                    the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        x = self.extract_feat(batch_inputs)
        results_list = self.bbox_head.predict(
            x, batch_data_samples, rescale=rescale)
        batch_data_samples = self.add_pred_to_datasample(
            batch_data_samples, results_list)
        return batch_data_samples

    def _execute(
            self,
            batch_inputs: jt.Var,
            batch_data_samples: OptSampleList = None) -> Tuple[List[jt.Var]]:
        """Network forward process. Usually includes backbone, neck and head
        forward without any post-processing.

         Args:
            batch_inputs (Tensor): Inputs with shape (N, C, H, W).

        Returns:
            tuple[list]: A tuple of features from ``bbox_head`` forward.
        """
        x = self.extract_feat(batch_inputs)
        results = self.bbox_head.execute(x)
        return results

    def extract_feat(self, batch_inputs: jt.Var) -> Tuple[jt.Var]:
        """Extract features.

        Args:
            batch_inputs (Tensor): Image tensor with shape (N, C, H ,W).

        Returns:
            tuple[Tensor]: Multi-level features that may have
            different resolutions.
        """
        x = self.backbone(batch_inputs)
        if self.with_neck:
            x = self.neck(x)
        return x

jittordet/models/frameworks/single_stage.py
JittorCV/jittordet
from .conv_module import ConvModule
from .linear import *  # noqa: F403, F401

__all__ = ['ConvModule']

jittordet/models/layers/__init__.py
JittorCV/jittordet
import warnings

from jittor import nn

from jittordet.engine import BRICKS
from jittordet.models.utils import constant_init, kaiming_init

# register padding bricks
BRICKS.register_module('zero', module=nn.ZeroPad2d)
BRICKS.register_module('reflect', module=nn.ReflectionPad2d)
BRICKS.register_module('replicate', module=nn.ReplicationPad2d)

# register conv bricks
BRICKS.register_module('Conv1d', module=nn.Conv1d)
BRICKS.register_module('Conv2d', module=nn.Conv2d)
BRICKS.register_module('Conv3d', module=nn.Conv3d)
BRICKS.register_module('Conv', module=nn.Conv2d)

# register
BRICKS.register_module('BN', module=nn.BatchNorm2d)
BRICKS.register_module('BN1d', module=nn.BatchNorm1d)
BRICKS.register_module('BN2d', module=nn.BatchNorm2d)
BRICKS.register_module('BN3d', module=nn.BatchNorm3d)
BRICKS.register_module('GN', module=nn.GroupNorm)
BRICKS.register_module('LN', module=nn.LayerNorm)
BRICKS.register_module('IN', module=nn.InstanceNorm2d)
BRICKS.register_module('IN1d', module=nn.InstanceNorm1d)
BRICKS.register_module('IN2d', module=nn.InstanceNorm2d)
BRICKS.register_module('IN3d', module=nn.InstanceNorm3d)

BRICKS.register_module('ReLU', module=nn.ReLU)
BRICKS.register_module(module=nn.LeakyReLU)
BRICKS.register_module(module=nn.PReLU)
BRICKS.register_module(module=nn.ReLU6)
BRICKS.register_module(module=nn.ELU)
BRICKS.register_module(module=nn.Sigmoid)
BRICKS.register_module(module=nn.Tanh)
BRICKS.register_module(module=nn.GELU)


class ConvModule(nn.Module):
    """A conv block that bundles conv/norm/activation layers.

    This block simplifies the usage of convolution layers, which are commonly
    used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
    It is based upon three build methods: `build_conv_layer()`,
    `build_norm_layer()` and `build_activation_layer()`.
    Besides, we add some additional features in this module.
    1. Automatically set `bias` of the conv layer.
    2. Spectral norm is supported.
    3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
    supports zero and circular padding, and we add "reflect" padding mode.
    Args:
        in_channels (int): Number of channels in the input feature map.
            Same as that in ``nn._ConvNd``.
        out_channels (int): Number of channels produced by the convolution.
            Same as that in ``nn._ConvNd``.
        kernel_size (int | tuple[int]): Size of the convolving kernel.
            Same as that in ``nn._ConvNd``.
        stride (int | tuple[int]): Stride of the convolution.
            Same as that in ``nn._ConvNd``.
        padding (int | tuple[int]): Zero-padding added to both sides of
            the input. Same as that in ``nn._ConvNd``.
        dilation (int | tuple[int]): Spacing between kernel elements.
            Same as that in ``nn._ConvNd``.
        groups (int): Number of blocked connections from input channels to
            output channels. Same as that in ``nn._ConvNd``.
        bias (bool | str): If specified as `auto`, it will be decided by the
            norm_cfg. Bias will be set as True if `norm_cfg` is None, otherwise
            False. Default: "auto".
        conv_cfg (dict): Config dict for convolution layer. Default: None,
            which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU').
        with_spectral_norm (bool): Whether use spectral norm in conv module.
            Default: False.
        padding_mode (str): If the `padding_mode` has not been supported by
            current `Conv2d` in PyTorch, we will use our own padding layer
            instead. Currently, we support ['zeros', 'circular'] with official
            implementation and ['reflect'] with our own implementation.
            Default: 'zeros'.
        order (tuple[str]): The order of conv/norm/activation layers. It is a
            sequence of "conv", "norm" and "act". Common examples are
            ("conv", "norm", "act") and ("act", "conv", "norm").
            Default: ('conv', 'norm', 'act').
    """

    _abbr_ = 'conv_block'

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias='auto',
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=dict(type='ReLU'),
                 with_spectral_norm=False,
                 padding_mode='zeros',
                 order=('conv', 'norm', 'act')):
        super(ConvModule, self).__init__()
        assert conv_cfg is None or isinstance(conv_cfg, dict)
        assert norm_cfg is None or isinstance(norm_cfg, dict)
        assert act_cfg is None or isinstance(act_cfg, dict)
        official_padding_mode = ['zeros', 'circular']
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.with_spectral_norm = with_spectral_norm
        self.with_explicit_padding = padding_mode not in official_padding_mode
        self.order = order
        assert isinstance(self.order, tuple) and len(self.order) == 3
        assert set(order) == set(['conv', 'norm', 'act'])

        self.with_norm = norm_cfg is not None
        self.with_activation = act_cfg is not None
        # if the conv layer is before a norm layer, bias is unnecessary.
        if bias == 'auto':
            bias = not self.with_norm
        self.with_bias = bias

        if self.with_norm and self.with_bias:
            warnings.warn('ConvModule has norm and bias at the same time')

        if self.with_explicit_padding:
            pad_cfg = dict(type=padding_mode)
            self.padding_layer = BRICKS.build(pad_cfg, padding=padding)

        # reset padding to 0 for conv module
        conv_padding = 0 if self.with_explicit_padding else padding
        # build convolution layer
        self.conv = BRICKS.build(
            dict(type='Conv2d') if conv_cfg is None else conv_cfg,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=conv_padding,
            dilation=dilation,
            groups=groups,
            bias=bias)
        # export the attributes of self.conv to a higher level for convenience
        self.in_channels = self.conv.in_channels
        self.out_channels = self.conv.out_channels
        self.kernel_size = self.conv.kernel_size
        self.stride = self.conv.stride
        self.padding = padding
        self.dilation = self.conv.dilation
        # self.transposed = self.conv.transposed
        # self.output_padding = self.conv.output_padding
        self.groups = self.conv.groups

        if self.with_spectral_norm:
            self.conv = nn.utils.spectral_norm(self.conv)

        # build normalization layers
        if self.with_norm:
            # norm layer is after conv layer
            if order.index('norm') > order.index('conv'):
                norm_channels = out_channels
            else:
                norm_channels = in_channels
            if norm_cfg.get('type', 'BN') == 'GN':
                self.gn = BRICKS.build(norm_cfg, num_channels=norm_channels)
                self.norm = 'gn'
            else:
                self.bn = BRICKS.build(norm_cfg, in_channels=norm_channels)
                self.norm = 'bn'
        else:
            self.norm = 'None'

        # build activation layer
        if self.with_activation:
            act_cfg_ = act_cfg.copy()
            if act_cfg_['type'] not in [
                    'Tanh', 'PReLU', 'Sigmoid', 'HSigmoid', 'Swish'
            ]:
                self.activate = BRICKS.build(act_cfg_)

        # Use msra init by default
        self.init_weights()

    def init_weights(self):
        # 1. It is mainly for customized conv layers with their own
        #    initialization manners by calling their own ``init_weights()``,
        #    and we do not want ConvModule to override the initialization.
        # 2. For customized conv layers without their own initialization
        #    manners (that is, they don't have their own ``init_weights()``)
        #    and PyTorch's conv layers, they will be initialized by
        #    this method with default ``kaiming_init``.
        # Note: For PyTorch's conv layers, they will be overwritten by our
        #    initialization implementation using default ``kaiming_init``.
        if not hasattr(self.conv, 'init_weights'):
            if self.with_activation and self.act_cfg['type'] == 'LeakyReLU':
                nonlinearity = 'leaky_relu'
                a = self.act_cfg.get('negative_slope', 0.01)
            else:
                nonlinearity = 'relu'
                a = 0
            kaiming_init(self.conv, a=a, nonlinearity=nonlinearity)
        if self.with_norm:
            constant_init(self.norm, 1, bias=0)

    def execute(self, x, activate=True, norm=True):
        for layer in self.order:
            if layer == 'conv':
                if self.with_explicit_padding:
                    x = self.padding_layer(x)
                x = self.conv(x)
            elif layer == 'norm' and norm and self.with_norm:
                x = getattr(self, self.norm)(x)
            elif layer == 'act' and activate and self.with_activation:
                x = self.activate(x)
        return x

jittordet/models/layers/conv_module.py
JittorCV/jittordet
import jittor.nn as nn

from jittordet.engine import MODELS

MODELS.register_module('Linear', module=nn.Linear)

jittordet/models/layers/linear.py
JittorCV/jittordet
from .accuracy import Accuracy, accuracy
from .cross_entropy_loss import CrossEntropyLoss
from .focal_loss import FocalLoss
from .iou_loss import GIoULoss, IoULoss
from .smooth_l1_loss import L1Loss, SmoothL1Loss
from .utils import reduce_loss, weight_reduce_loss, weighted_loss

__all__ = [
    'reduce_loss', 'weight_reduce_loss', 'weighted_loss', 'CrossEntropyLoss',
    'Accuracy', 'accuracy', 'FocalLoss', 'IoULoss', 'GIoULoss', 'SmoothL1Loss',
    'L1Loss'
]

jittordet/models/losses/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/losses/accuracy.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS


def accuracy(pred, target, topk=1, thresh=None):
    """Calculate accuracy according to the prediction and target.

    Args:
        pred (torch.Tensor): The model prediction, shape (N, num_class)
        target (torch.Tensor): The target of each prediction, shape (N, )
        topk (int | tuple[int], optional): If the predictions in ``topk``
            matches the target, the predictions will be regarded as
            correct ones. Defaults to 1.
        thresh (float, optional): If not None, predictions with scores under
            this threshold are considered incorrect. Default to None.

    Returns:
        float | tuple[float]: If the input ``topk`` is a single integer,
            the function will return a single float as accuracy. If
            ``topk`` is a tuple containing multiple integers, the
            function will return a tuple containing accuracies of
            each ``topk`` number.
    """
    assert isinstance(topk, (int, tuple))
    if isinstance(topk, int):
        topk = (topk, )
        return_single = True
    else:
        return_single = False

    maxk = max(topk)
    if pred.size(0) == 0:
        accu = [jt.array(0, dtype=pred.dtype) for i in range(len(topk))]
        return accu[0] if return_single else accu
    assert pred.ndim == 2 and target.ndim == 1
    assert pred.size(0) == target.size(0)
    assert maxk <= pred.size(1), \
        f'maxk {maxk} exceeds pred dimension {pred.size(1)}'
    pred_value, pred_label = pred.topk(maxk, dim=1)
    pred_label = pred_label.t()  # transpose to shape (maxk, N)

    correct = jt.equal(pred_label, target.view(1, -1).expand_as(pred_label))
    if thresh is not None:
        # Only prediction values larger than thresh are counted as correct
        correct = correct & (pred_value > thresh).t()
    res = []
    for k in topk:
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdims=True)
        res.append(jt.multiply(correct_k, 100.0 / pred.size(0)))
    return res[0] if return_single else res


@MODELS.register_module()
class Accuracy(nn.Module):

    def __init__(self, topk=(1, ), thresh=None):
        """Module to calculate the accuracy.

        Args:
            topk (tuple, optional): The criterion used to calculate the
                accuracy. Defaults to (1,).
            thresh (float, optional): If not None, predictions with scores
                under this threshold are considered incorrect. Default to None.
        """
        super().__init__()
        self.topk = topk
        self.thresh = thresh

    def execute(self, pred, target):
        """Forward function to calculate accuracy.

        Args:
            pred (torch.Tensor): Prediction of models.
            target (torch.Tensor): Target for each prediction.

        Returns:
            tuple[float]: The accuracies under different topk criterions.
        """
        return accuracy(pred, target, self.topk, self.thresh)

jittordet/models/losses/accuracy.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/losses/cross_entropy_loss.py
# Copyright (c) OpenMMLab. All rights reserved.
import warnings

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS
from .utils import weight_reduce_loss


def binary_cross_entropy_with_logits(output,
                                     target,
                                     weight=None,
                                     pos_weight=None,
                                     reduction='mean'):
    max_val = jt.clamp(-output, min_v=0)
    if pos_weight is not None:
        log_weight = (pos_weight - 1) * target + 1
        loss = (1 - target) * output + (
            log_weight * (((-max_val).exp() +
                           (-output - max_val).exp()).log() + max_val))
    else:
        loss = (1 - target) * output + max_val + (
            (-max_val).exp() + (-output - max_val).exp()).log()
    if weight is not None:
        loss *= weight

    if reduction == 'mean':
        return loss.mean()
    elif reduction == 'sum':
        return loss.sum()
    elif reduction == 'none':
        return loss
    else:
        raise RuntimeError(f'Got invalid reduction, {reduction}')


def cross_entropy(pred,
                  label,
                  weight=None,
                  reduction='mean',
                  avg_factor=None,
                  class_weight=None,
                  ignore_index=-100,
                  avg_non_ignore=False):
    """Calculate the CrossEntropy loss.

    Args:
        pred (torch.Tensor): The prediction with shape (N, C), C is the number
            of classes.
        label (torch.Tensor): The learning label of the prediction.
        weight (torch.Tensor, optional): Sample-wise loss weight.
        reduction (str, optional): The method used to reduce the loss.
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
        class_weight (list[float], optional): The weight for each class.
        ignore_index (int | None): The label index to be ignored.
            If None, it will be set to default value. Default: -100.
        avg_non_ignore (bool): The flag decides to whether the loss is
            only averaged over non-ignored targets. Default: False.

    Returns:
        torch.Tensor: The calculated loss
    """
    # The default value of ignore_index is the same as F.cross_entropy
    ignore_index = -100 if ignore_index is None else ignore_index
    # element-wise losses
    loss = nn.cross_entropy_loss(
        pred,
        label,
        weight=class_weight,
        reduction='none',
        ignore_index=ignore_index)

    # average loss over non-ignored elements
    # pytorch's official cross_entropy average loss over non-ignored elements
    # refer to https://github.com/pytorch/pytorch/blob/56b43f4fec1f76953f15a627694d4bba34588969/torch/nn/functional.py#L2660  # noqa
    if (avg_factor is None) and avg_non_ignore and reduction == 'mean':
        avg_factor = label.numel() - (label == ignore_index).sum().item()

    # apply weights and do the reduction
    if weight is not None:
        weight = weight.float()
    loss = weight_reduce_loss(
        loss, weight=weight, reduction=reduction, avg_factor=avg_factor)

    return loss


def _expand_onehot_labels(labels, label_weights, label_channels, ignore_index):
    """Expand onehot labels to match the size of prediction."""
    bin_labels = jt.full((labels.size(0), label_channels),
                         0,
                         dtype=labels.dtype)
    valid_mask = (labels >= 0) & (labels != ignore_index)
    inds = jt.nonzero(valid_mask & (labels < label_channels))

    if inds.numel() > 0:
        bin_labels[inds, labels[inds]] = 1

    valid_mask = valid_mask.view(-1, 1).expand(labels.size(0),
                                               label_channels).float()
    if label_weights is None:
        bin_label_weights = valid_mask
    else:
        bin_label_weights = label_weights.view(-1, 1).repeat(1, label_channels)
        bin_label_weights *= valid_mask

    return bin_labels, bin_label_weights, valid_mask


def binary_cross_entropy(pred,
                         label,
                         weight=None,
                         reduction='mean',
                         avg_factor=None,
                         class_weight=None,
                         ignore_index=-100,
                         avg_non_ignore=False):
    """Calculate the binary CrossEntropy loss.

    Args:
        pred (torch.Tensor): The prediction with shape (N, 1) or (N, ).
            When the shape of pred is (N, 1), label will be expanded to
            one-hot format, and when the shape of pred is (N, ), label
            will not be expanded to one-hot format.
        label (torch.Tensor): The learning label of the prediction,
            with shape (N, ).
        weight (torch.Tensor, optional): Sample-wise loss weight.
        reduction (str, optional): The method used to reduce the loss.
            Options are "none", "mean" and "sum".
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
        class_weight (list[float], optional): The weight for each class.
        ignore_index (int | None): The label index to be ignored.
            If None, it will be set to default value. Default: -100.
        avg_non_ignore (bool): The flag decides to whether the loss is
            only averaged over non-ignored targets. Default: False.

    Returns:
        torch.Tensor: The calculated loss.
    """
    # The default value of ignore_index is the same as F.cross_entropy
    ignore_index = -100 if ignore_index is None else ignore_index

    if pred.ndim != label.ndim:
        label, weight, valid_mask = _expand_onehot_labels(
            label, weight, pred.size(-1), ignore_index)
    else:
        # should mask out the ignored elements
        valid_mask = ((label >= 0) & (label != ignore_index)).float()
        if weight is not None:
            # The inplace writing method will have a mismatched broadcast
            # shape error if the weight and valid_mask dimensions
            # are inconsistent such as (B,N,1) and (B,N,C).
            weight = weight * valid_mask
        else:
            weight = valid_mask

    # average loss over non-ignored elements
    if (avg_factor is None) and avg_non_ignore and reduction == 'mean':
        avg_factor = valid_mask.sum().item()

    # weighted element-wise losses
    weight = weight.float()
    loss = binary_cross_entropy_with_logits(
        pred, label.float(), pos_weight=class_weight, reduction='none')
    # do the reduction for the weighted loss
    loss = weight_reduce_loss(
        loss, weight, reduction=reduction, avg_factor=avg_factor)

    return loss


def mask_cross_entropy(pred,
                       target,
                       label,
                       reduction='mean',
                       avg_factor=None,
                       class_weight=None,
                       ignore_index=None,
                       **kwargs):
    """Calculate the CrossEntropy loss for masks.

    Args:
        pred (torch.Tensor): The prediction with shape (N, C, *), C is the
            number of classes. The trailing * indicates arbitrary shape.
        target (torch.Tensor): The learning label of the prediction.
        label (torch.Tensor): ``label`` indicates the class label of the mask
            corresponding object. This will be used to select the mask in the
            of the class which the object belongs to when the mask prediction
            if not class-agnostic.
        reduction (str, optional): The method used to reduce the loss.
            Options are "none", "mean" and "sum".
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
        class_weight (list[float], optional): The weight for each class.
        ignore_index (None): Placeholder, to be consistent with other loss.
            Default: None.

    Returns:
        torch.Tensor: The calculated loss

    Example:
        >>> N, C = 3, 11
        >>> H, W = 2, 2
        >>> pred = torch.randn(N, C, H, W) * 1000
        >>> target = torch.rand(N, H, W)
        >>> label = torch.randint(0, C, size=(N,))
        >>> reduction = 'mean'
        >>> avg_factor = None
        >>> class_weights = None
        >>> loss = mask_cross_entropy(pred, target, label, reduction,
        >>>                           avg_factor, class_weights)
        >>> assert loss.shape == (1,)
    """
    assert ignore_index is None, 'BCE loss does not support ignore_index'
    # TODO: handle these two reserved arguments
    assert reduction == 'mean' and avg_factor is None
    num_rois = pred.size()[0]
    inds = jt.arange(0, num_rois, dtype=jt.int64)
    pred_slice = pred[inds, label].squeeze(1)
    return binary_cross_entropy_with_logits(
        pred_slice, target, weight=class_weight, reduction='mean')[None]


@MODELS.register_module()
class CrossEntropyLoss(nn.Module):

    def __init__(self,
                 use_sigmoid=False,
                 use_mask=False,
                 reduction='mean',
                 class_weight=None,
                 ignore_index=None,
                 loss_weight=1.0,
                 avg_non_ignore=False):
        """CrossEntropyLoss.

        Args:
            use_sigmoid (bool, optional): Whether the prediction uses sigmoid
                of softmax. Defaults to False.
            use_mask (bool, optional): Whether to use mask cross entropy loss.
                Defaults to False.
            reduction (str, optional): . Defaults to 'mean'.
                Options are "none", "mean" and "sum".
            class_weight (list[float], optional): Weight of each class.
                Defaults to None.
            ignore_index (int | None): The label index to be ignored.
                Defaults to None.
            loss_weight (float, optional): Weight of the loss. Defaults to 1.0.
            avg_non_ignore (bool): The flag decides to whether the loss is
                only averaged over non-ignored targets. Default: False.
        """
        super(CrossEntropyLoss, self).__init__()
        assert (use_sigmoid is False) or (use_mask is False)
        self.use_sigmoid = use_sigmoid
        self.use_mask = use_mask
        self.reduction = reduction
        self.loss_weight = loss_weight
        self.class_weight = class_weight
        self.ignore_index = ignore_index
        self.avg_non_ignore = avg_non_ignore
        if ((ignore_index is not None) and not self.avg_non_ignore
                and self.reduction == 'mean'):
            warnings.warn(
                'Default ``avg_non_ignore`` is False, if you would like to '
                'ignore the certain label and average loss over non-ignore '
                'labels, which is the same with PyTorch official '
                'cross_entropy, set ``avg_non_ignore=True``.')

        if self.use_sigmoid:
            self.cls_criterion = binary_cross_entropy
        elif self.use_mask:
            self.cls_criterion = mask_cross_entropy
        else:
            self.cls_criterion = cross_entropy

    def extra_repr(self):
        """Extra repr."""
        s = f'avg_non_ignore={self.avg_non_ignore}'
        return s

    def execute(self,
                cls_score,
                label,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                ignore_index=None,
                **kwargs):
        """Forward function.

        Args:
            cls_score (torch.Tensor): The prediction.
            label (torch.Tensor): The learning label of the prediction.
            weight (torch.Tensor, optional): Sample-wise loss weight.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The method used to reduce the
                loss. Options are "none", "mean" and "sum".
            ignore_index (int | None): The label index to be ignored.
                If not None, it will override the default value. Default: None.
        Returns:
            torch.Tensor: The calculated loss.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        if ignore_index is None:
            ignore_index = self.ignore_index

        if self.class_weight is not None:
            class_weight = jt.array(self.class_weight, dtype=cls_score.dtype)
        else:
            class_weight = None
        loss_cls = self.loss_weight * self.cls_criterion(
            cls_score,
            label,
            weight,
            class_weight=class_weight,
            reduction=reduction,
            avg_factor=avg_factor,
            ignore_index=ignore_index,
            avg_non_ignore=self.avg_non_ignore,
            **kwargs)
        return loss_cls

jittordet/models/losses/cross_entropy_loss.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/losses/focal_loss.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor.nn as nn

from jittordet.engine import MODELS
from .cross_entropy_loss import binary_cross_entropy_with_logits
from .utils import weight_reduce_loss


# This method is only for debugging
def py_sigmoid_focal_loss(pred,
                          target,
                          weight=None,
                          gamma=2.0,
                          alpha=0.25,
                          reduction='mean',
                          avg_factor=None):
    """PyTorch version of `Focal Loss <https://arxiv.org/abs/1708.02002>`_.

    Args:
        pred (torch.Tensor): The prediction with shape (N, C), C is the
            number of classes
        target (torch.Tensor): The learning label of the prediction.
        weight (torch.Tensor, optional): Sample-wise loss weight.
        gamma (float, optional): The gamma for calculating the modulating
            factor. Defaults to 2.0.
        alpha (float, optional): A balanced form for Focal Loss.
            Defaults to 0.25.
        reduction (str, optional): The method used to reduce the loss into
            a scalar. Defaults to 'mean'.
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
    """
    pred_sigmoid = pred.sigmoid()
    target = target.astype(pred.dtype)
    pt = (1 - pred_sigmoid) * target + pred_sigmoid * (1 - target)
    focal_weight = (alpha * target + (1 - alpha) *
                    (1 - target)) * pt.pow(gamma)
    loss = binary_cross_entropy_with_logits(
        pred, target, reduction='none') * focal_weight
    if weight is not None:
        if weight.shape != loss.shape:
            if weight.size(0) == loss.size(0):
                # For most cases, weight is of shape (num_priors, ),
                #  which means it does not have the second axis num_class
                weight = weight.view(-1, 1)
            else:
                # Sometimes, weight per anchor per class is also needed. e.g.
                #  in FSAF. But it may be flattened of shape
                #  (num_priors x num_class, ), while loss is still of shape
                #  (num_priors, num_class).
                assert weight.numel() == loss.numel()
                weight = weight.view(loss.size(0), -1)
        assert weight.ndim == loss.ndim
    loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
    return loss


@MODELS.register_module()
class FocalLoss(nn.Module):

    def __init__(self,
                 use_sigmoid=True,
                 gamma=2.0,
                 alpha=0.25,
                 reduction='mean',
                 loss_weight=1.0):
        """`Focal Loss <https://arxiv.org/abs/1708.02002>`_

        Args:
            use_sigmoid (bool, optional): Whether to the prediction is
                used for sigmoid or softmax. Defaults to True.
            gamma (float, optional): The gamma for calculating the modulating
                factor. Defaults to 2.0.
            alpha (float, optional): A balanced form for Focal Loss.
                Defaults to 0.25.
            reduction (str, optional): The method used to reduce the loss into
                a scalar. Defaults to 'mean'. Options are "none", "mean" and
                "sum".
            loss_weight (float, optional): Weight of loss. Defaults to 1.0.
            activated (bool, optional): Whether the input is activated.
                If True, it means the input has been activated and can be
                treated as probabilities. Else, it should be treated as logits.
                Defaults to False.
        """
        super(FocalLoss, self).__init__()
        assert use_sigmoid is True, 'Only sigmoid focal loss supported now.'
        self.use_sigmoid = use_sigmoid
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning label of the prediction.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Options are "none", "mean" and "sum".

        Returns:
            torch.Tensor: The calculated loss
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        if self.use_sigmoid:
            num_classes = pred.size(1)
            target = nn.one_hot(target, num_classes=num_classes + 1)
            target = target[:, :num_classes]
            loss_cls = self.loss_weight * py_sigmoid_focal_loss(
                pred,
                target,
                weight,
                gamma=self.gamma,
                alpha=self.alpha,
                reduction=reduction,
                avg_factor=avg_factor)
        else:
            raise NotImplementedError
        return loss_cls

jittordet/models/losses/focal_loss.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
import warnings

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS
from jittordet.utils import bbox_overlaps
from .utils import weighted_loss


@weighted_loss
def iou_loss(pred, target, linear=False, mode='log', eps=1e-6):
    """IoU loss.

    Computing the IoU loss between a set of predicted bboxes and target bboxes.
    The loss is calculated as negative log of IoU.

    Args:
        pred (torch.Tensor): Predicted bboxes of format (x1, y1, x2, y2),
            shape (n, 4).
        target (torch.Tensor): Corresponding gt bboxes, shape (n, 4).
        linear (bool, optional): If True, use linear scale of loss instead of
            log scale. Default: False.
        mode (str): Loss scaling mode, including "linear", "square", and "log".
            Default: 'log'
        eps (float): Eps to avoid log(0).

    Return:
        torch.Tensor: Loss tensor.
    """
    assert mode in ['linear', 'square', 'log']
    if linear:
        mode = 'linear'
        warnings.warn('DeprecationWarning: Setting "linear=True" in '
                      'iou_loss is deprecated, please use "mode=`linear`" '
                      'instead.')
    ious = bbox_overlaps(pred, target, is_aligned=True).clamp(min_v=eps)
    if mode == 'linear':
        loss = 1 - ious
    elif mode == 'square':
        loss = 1 - ious**2
    elif mode == 'log':
        loss = -ious.log()
    else:
        raise NotImplementedError
    return loss


@weighted_loss
def giou_loss(pred, target, eps=1e-7):
    r"""`Generalized Intersection over Union: A Metric and A Loss for Bounding
    Box Regression <https://arxiv.org/abs/1902.09630>`_.

    Args:
        pred (torch.Tensor): Predicted bboxes of format (x1, y1, x2, y2),
            shape (n, 4).
        target (torch.Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).

    Return:
        Tensor: Loss tensor.
    """
    gious = bbox_overlaps(pred, target, mode='giou', is_aligned=True, eps=eps)
    loss = 1 - gious
    return loss


@weighted_loss
def diou_loss(pred, target, eps=1e-7):
    r"""`Implementation of Distance-IoU Loss: Faster and Better
    Learning for Bounding Box Regression, https://arxiv.org/abs/1911.08287`_.

    Code is modified from https://github.com/Zzh-tju/DIoU.

    Args:
        pred (Tensor): Predicted bboxes of format (x1, y1, x2, y2),
            shape (n, 4).
        target (Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).
    Return:
        Tensor: Loss tensor.
    """
    # overlap
    lt = jt.maximum(pred[:, :2], target[:, :2])
    rb = jt.minimum(pred[:, 2:], target[:, 2:])
    wh = (rb - lt).clamp(min=0)
    overlap = wh[:, 0] * wh[:, 1]

    # union
    ap = (pred[:, 2] - pred[:, 0]) * (pred[:, 3] - pred[:, 1])
    ag = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
    union = ap + ag - overlap + eps

    # IoU
    ious = overlap / union

    # enclose area
    enclose_x1y1 = jt.minimum(pred[:, :2], target[:, :2])
    enclose_x2y2 = jt.maximum(pred[:, 2:], target[:, 2:])
    enclose_wh = (enclose_x2y2 - enclose_x1y1).clamp(min=0)

    cw = enclose_wh[:, 0]
    ch = enclose_wh[:, 1]

    c2 = cw**2 + ch**2 + eps

    b1_x1, b1_y1 = pred[:, 0], pred[:, 1]
    b1_x2, b1_y2 = pred[:, 2], pred[:, 3]
    b2_x1, b2_y1 = target[:, 0], target[:, 1]
    b2_x2, b2_y2 = target[:, 2], target[:, 3]

    left = ((b2_x1 + b2_x2) - (b1_x1 + b1_x2))**2 / 4
    right = ((b2_y1 + b2_y2) - (b1_y1 + b1_y2))**2 / 4
    rho2 = left + right

    # DIoU
    dious = ious - rho2 / c2
    loss = 1 - dious
    return loss


@MODELS.register_module()
class IoULoss(nn.Module):
    """IoULoss.

    Computing the IoU loss between a set of predicted bboxes and target bboxes.

    Args:
        linear (bool): If True, use linear scale of loss else determined
            by mode. Default: False.
        eps (float): Eps to avoid log(0).
        reduction (str): Options are "none", "mean" and "sum".
        loss_weight (float): Weight of loss.
        mode (str): Loss scaling mode, including "linear", "square", and "log".
            Default: 'log'
    """

    def __init__(self,
                 linear=False,
                 eps=1e-6,
                 reduction='mean',
                 loss_weight=1.0,
                 mode='log'):
        super(IoULoss, self).__init__()
        assert mode in ['linear', 'square', 'log']
        if linear:
            mode = 'linear'
            warnings.warn('DeprecationWarning: Setting "linear=True" in '
                          'IOULoss is deprecated, please use "mode=`linear`" '
                          'instead.')
        self.mode = mode
        self.linear = linear
        self.eps = eps
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                **kwargs):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None. Options are "none", "mean" and "sum".
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        if (weight is not None) and (not jt.any(weight > 0)) and (reduction !=
                                                                  'none'):
            if pred.ndim == weight.ndim + 1:
                weight = weight.unsqueeze(1)
            return (pred * weight).sum()  # 0
        if weight is not None and weight.ndim > 1:
            # TODO: remove this in the future
            # reduce the weight of shape (n, 4) to (n,) to match the
            # iou_loss of shape (n,)
            assert weight.shape == pred.shape
            weight = weight.mean(-1)
        loss = self.loss_weight * iou_loss(
            pred,
            target,
            weight,
            mode=self.mode,
            eps=self.eps,
            reduction=reduction,
            avg_factor=avg_factor,
            **kwargs)
        return loss


@MODELS.register_module()
class GIoULoss(nn.Module):

    def __init__(self, eps=1e-6, reduction='mean', loss_weight=1.0):
        super(GIoULoss, self).__init__()
        self.eps = eps
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                **kwargs):
        if weight is not None and not jt.any(weight > 0):
            if pred.ndim == weight.ndim + 1:
                weight = weight.unsqueeze(1)
            return (pred * weight).sum()  # 0
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        if weight is not None and weight.ndim > 1:
            # TODO: remove this in the future
            # reduce the weight of shape (n, 4) to (n,) to match the
            # giou_loss of shape (n,)
            assert weight.shape == pred.shape
            weight = weight.mean(-1)
        loss = self.loss_weight * giou_loss(
            pred,
            target,
            weight,
            eps=self.eps,
            reduction=reduction,
            avg_factor=avg_factor,
            **kwargs)
        return loss


@MODELS.register_module()
class DIoULoss(nn.Module):

    def __init__(self, eps=1e-6, reduction='mean', loss_weight=1.0):
        super(DIoULoss, self).__init__()
        self.eps = eps
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                **kwargs):
        if weight is not None and not jt.any(weight > 0):
            if pred.ndim == weight.ndim + 1:
                weight = weight.unsqueeze(1)
            return (pred * weight).sum()  # 0
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        if weight is not None and weight.ndim > 1:
            # TODO: remove this in the future
            # reduce the weight of shape (n, 4) to (n,) to match the
            # giou_loss of shape (n,)
            assert weight.shape == pred.shape
            weight = weight.mean(-1)
        loss = self.loss_weight * diou_loss(
            pred,
            target,
            weight,
            eps=self.eps,
            reduction=reduction,
            avg_factor=avg_factor,
            **kwargs)
        return loss

jittordet/models/losses/iou_loss.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/losses/smooth_l1_loss.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS
from .utils import weighted_loss


@weighted_loss
def smooth_l1_loss(pred, target, beta=1.0):
    """Smooth L1 loss.

    Args:
        pred (torch.Tensor): The prediction.
        target (torch.Tensor): The learning target of the prediction.
        beta (float, optional): The threshold in the piecewise function.
            Defaults to 1.0.

    Returns:
        torch.Tensor: Calculated loss
    """
    assert beta > 0
    if target.numel() == 0:
        return pred.sum() * 0

    assert pred.size() == target.size()
    diff = jt.abs(pred - target)
    condition = (diff < beta).astype(diff.dtype)
    loss = (0.5 * diff * diff / beta) * condition + \
        (diff - 0.5 * beta) * (1 - condition)
    return loss


@weighted_loss
def l1_loss(pred, target):
    """L1 loss.

    Args:
        pred (torch.Tensor): The prediction.
        target (torch.Tensor): The learning target of the prediction.

    Returns:
        torch.Tensor: Calculated loss
    """
    if target.numel() == 0:
        return pred.sum() * 0

    assert pred.size() == target.size()
    loss = jt.abs(pred - target)
    return loss


@MODELS.register_module()
class SmoothL1Loss(nn.Module):
    """Smooth L1 loss.

    Args:
        beta (float, optional): The threshold in the piecewise function.
            Defaults to 1.0.
        reduction (str, optional): The method to reduce the loss.
            Options are "none", "mean" and "sum". Defaults to "mean".
        loss_weight (float, optional): The weight of loss.
    """

    def __init__(self, beta=1.0, reduction='mean', loss_weight=1.0):
        super(SmoothL1Loss, self).__init__()
        self.beta = beta
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                **kwargs):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss_bbox = self.loss_weight * smooth_l1_loss(
            pred,
            target,
            weight,
            beta=self.beta,
            reduction=reduction,
            avg_factor=avg_factor,
            **kwargs)
        return loss_bbox


@MODELS.register_module()
class L1Loss(nn.Module):
    """L1 loss.

    Args:
        reduction (str, optional): The method to reduce the loss.
            Options are "none", "mean" and "sum".
        loss_weight (float, optional): The weight of loss.
    """

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(L1Loss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def execute(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        """Forward function.

        Args:
            pred (torch.Tensor): The prediction.
            target (torch.Tensor): The learning target of the prediction.
            weight (torch.Tensor, optional): The weight of loss for each
                prediction. Defaults to None.
            avg_factor (int, optional): Average factor that is used to average
                the loss. Defaults to None.
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None.
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss_bbox = self.loss_weight * l1_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss_bbox

jittordet/models/losses/smooth_l1_loss.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/losses/utils.py
# Copyright (c) OpenMMLab. All rights reserved.
import functools


def reduce_loss(loss, reduction):
    """Reduce loss as specified.

    Args:
        loss (Tensor): Elementwise loss tensor.
        reduction (str): Options are "none", "mean" and "sum".

    Return:
        Tensor: Reduced loss tensor.
    """
    # none: 0, elementwise_mean:1, sum: 2
    if reduction == 'none':
        return loss
    elif reduction == 'mean':
        return loss.mean()
    elif reduction == 'sum':
        return loss.sum()


def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None):
    """Apply element-wise weight and reduce loss.

    Args:
        loss (Tensor): Element-wise loss.
        weight (Tensor): Element-wise weights.
        reduction (str): Same as built-in losses of PyTorch.
        avg_factor (float): Average factor when computing the mean of losses.

    Returns:
        Tensor: Processed loss values.
    """
    # if weight is specified, apply element-wise weight
    if weight is not None:
        loss = loss * weight

    # if avg_factor is not specified, just reduce the loss
    if avg_factor is None:
        loss = reduce_loss(loss, reduction)
    else:
        # if reduction is mean, then average the loss by avg_factor
        if reduction == 'mean':
            # Avoid causing ZeroDivisionError when avg_factor is 0.0,
            # i.e., all labels of an image belong to ignore index.
            eps = 1e-6
            loss = loss.sum() / (avg_factor + eps)
        # if reduction is 'none', then do nothing, otherwise raise an error
        elif reduction != 'none':
            raise ValueError('avg_factor can not be used with reduction="sum"')
    return loss


def weighted_loss(loss_func):
    """Create a weighted version of a given loss function.

    To use this decorator, the loss function must have the signature like
    `loss_func(pred, target, **kwargs)`. The function only needs to compute
    element-wise loss without any reduction. This decorator will add weight
    and reduction arguments to the function. The decorated function will have
    the signature like `loss_func(pred, target, weight=None, reduction='mean',
    avg_factor=None, **kwargs)`.

    :Example:

    >>> import torch
    >>> @weighted_loss
    >>> def l1_loss(pred, target):
    >>>     return (pred - target).abs()

    >>> pred = torch.Tensor([0, 2, 3])
    >>> target = torch.Tensor([1, 1, 1])
    >>> weight = torch.Tensor([1, 0, 1])

    >>> l1_loss(pred, target)
    tensor(1.3333)
    >>> l1_loss(pred, target, weight)
    tensor(1.)
    >>> l1_loss(pred, target, reduction='none')
    tensor([1., 1., 2.])
    >>> l1_loss(pred, target, weight, avg_factor=2)
    tensor(1.5000)
    """

    @functools.wraps(loss_func)
    def wrapper(pred,
                target,
                weight=None,
                reduction='mean',
                avg_factor=None,
                **kwargs):
        # get element-wise loss
        loss = loss_func(pred, target, **kwargs)
        loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss

    return wrapper

jittordet/models/losses/utils.py
JittorCV/jittordet
from jittor import nn

from jittordet.engine import MODELS
from jittordet.models.layers import ConvModule
from jittordet.models.utils import xavier_init


@MODELS.register_module()
class FPN(nn.Module):
    r"""Feature Pyramid Network.
    This is an implementation of paper `Feature Pyramid Networks for Object
    Detection <https://arxiv.org/abs/1612.03144>`_.
    Args:
        in_channels (List[int]): Number of input channels per scale.
        out_channels (int): Number of output channels (used at each scale)
        num_outs (int): Number of output scales.
        start_level (int): Index of the start input backbone level used to
            build the feature pyramid. Default: 0.
        end_level (int): Index of the end input backbone level (exclusive) to
            build the feature pyramid. Default: -1, which means the last level.
        add_extra_convs (bool | str): If bool, it decides whether to add conv
            layers on top of the original feature maps. Default to False.
            If True, its actual mode is specified by `extra_convs_on_inputs`.
            If str, it specifies the source feature map of the extra convs.
            Only the following options are allowed
            - 'on_input': Last feat map of neck inputs (i.e. backbone feature).
            - 'on_lateral':  Last feature map after lateral convs.
            - 'on_output': The last output feature map after fpn convs.
        extra_convs_on_inputs (bool, deprecated): Whether to apply extra convs
            on the original feature from the backbone. If True,
            it is equivalent to `add_extra_convs='on_input'`. If False, it is
            equivalent to set `add_extra_convs='on_output'`. Default to True.
        relu_before_extra_convs (bool): Whether to apply relu before the extra
            conv. Default: False.
        no_norm_on_lateral (bool): Whether to apply norm on lateral.
            Default: False.
        conv_cfg (dict): Config dict for convolution layer. Default: None.
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (str): Config dict for activation layer in ConvModule.
            Default: None.
        upsample_cfg (dict): Config dict for interpolate layer.
            Default: `dict(mode='nearest')`
        init_cfg (dict or list[dict], optional): Initialization config dict.
    Example:
        >>> in_channels = [2, 3, 5, 7]
        >>> scales = [340, 170, 84, 43]
        >>> inputs = [jt.randn(1, c, s, s)
        ...           for c, s in zip(in_channels, scales)]
        >>> self = FPN(in_channels, 11, len(in_channels)).eval()
        >>> outputs = self.forward(inputs)
        >>> for i in range(len(outputs)):
        ...     print(f'outputs[{i}].shape = {outputs[i].shape}')
        outputs[0].shape = ([1, 11, 340, 340])
        outputs[1].shape = ([1, 11, 170, 170])
        outputs[2].shape = ([1, 11, 84, 84])
        outputs[3].shape = ([1, 11, 43, 43])
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 num_outs,
                 start_level=0,
                 end_level=-1,
                 add_extra_convs=False,
                 relu_before_extra_convs=False,
                 no_norm_on_lateral=False,
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=None,
                 upsample_cfg=dict(mode='nearest'),
                 upsample_div_factor=1):
        super().__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.relu_before_extra_convs = relu_before_extra_convs
        self.no_norm_on_lateral = no_norm_on_lateral
        self.upsample_cfg = upsample_cfg.copy()
        self.upsample_div_factor = upsample_div_factor
        if end_level == -1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level < inputs, no extra level is allowed
            self.backbone_end_level = end_level
            assert end_level <= len(in_channels)
            assert num_outs == end_level - start_level
        self.start_level = start_level
        self.end_level = end_level
        self.add_extra_convs = add_extra_convs
        assert isinstance(add_extra_convs, (str, bool))
        if isinstance(add_extra_convs, str):
            assert add_extra_convs in ('on_input', 'on_lateral', 'on_output')

        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()

        for i in range(self.start_level, self.backbone_end_level):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,
                act_cfg=act_cfg)
            fpn_conv = ConvModule(
                out_channels,
                out_channels,
                3,
                padding=1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)

            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)

        # add extra conv layers (e.g., RetinaNet)
        extra_levels = num_outs - self.backbone_end_level + self.start_level
        if self.add_extra_convs and extra_levels >= 1:
            for i in range(extra_levels):
                if i == 0 and self.add_extra_convs == 'on_input':
                    in_channels = self.in_channels[self.backbone_end_level - 1]
                else:
                    in_channels = out_channels
                extra_fpn_conv = ConvModule(
                    in_channels,
                    out_channels,
                    3,
                    stride=2,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg)
                self.fpn_convs.append(extra_fpn_conv)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution='uniform')

    def execute(self, inputs):
        """Forward function."""
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                laterals[i - 1] += nn.interpolate(laterals[i],
                                                  **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] += nn.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)
            laterals[i - 1] /= self.upsample_div_factor
        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(nn.pool(outs[-1], 1, stride=2, op='maximum'))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](nn.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)

jittordet/models/necks/fpn.py
JittorCV/jittordet
from .preprocessor import Preprocessor

__all__ = ['Preprocessor']

jittordet/models/preprocessors/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/data_preprocessors/data_preprocessor.py
# Copyright (c) OpenMMLab. All rights reserved.
import math
from numbers import Number
from typing import List, Optional, Sequence, Union

import jittor as jt
import jittor.nn as nn
import numpy as np

from jittordet.engine import MODELS
from jittordet.structures import DetDataSample
from jittordet.utils import is_list_of


@MODELS.register_module()
class Preprocessor(nn.Module):
    """Image pre-processor for detection tasks.

    Comparing with the :class:`mmengine.ImgDataPreprocessor`,

    1. It supports batch augmentations.
    2. It will additionally append batch_input_shape and pad_shape
    to data_samples considering the object detection task.

    It provides the data pre-processing as follows

    - Collate and move data to the target device.
    - Pad inputs to the maximum size of current batch with defined
      ``pad_value``. The padding size can be divisible by a defined
      ``pad_size_divisor``
    - Stack inputs to batch_inputs.
    - Convert inputs from bgr to rgb if the shape of input is (3, H, W).
    - Normalize image with defined std and mean.
    - Do batch augmentations during training.

    Args:
        mean (Sequence[Number], optional): The pixel mean of R, G, B channels.
            Defaults to None.
        std (Sequence[Number], optional): The pixel standard deviation of
            R, G, B channels. Defaults to None.
        pad_size_divisor (int): The size of padded image should be
            divisible by ``pad_size_divisor``. Defaults to 1.
        pad_value (Number): The padded pixel value. Defaults to 0.
        pad_mask (bool): Whether to pad instance masks. Defaults to False.
        mask_pad_value (int): The padded pixel value for instance masks.
            Defaults to 0.
        pad_seg (bool): Whether to pad semantic segmentation maps.
            Defaults to False.
        seg_pad_value (int): The padded pixel value for semantic
            segmentation maps. Defaults to 255.
        bgr_to_rgb (bool): whether to convert image from BGR to RGB.
            Defaults to False.
        rgb_to_bgr (bool): whether to convert image from RGB to RGB.
            Defaults to False.
        boxtype2tensor (bool): Whether to keep the ``BaseBoxes`` type of
            bboxes data or not. Defaults to False.
        batch_augments (list[dict], optional): Batch-level augmentations
    """

    def __init__(self,
                 mean: Sequence[Number] = None,
                 std: Sequence[Number] = None,
                 pad_size_divisor: int = 1,
                 pad_value: Union[float, int] = 0,
                 pad_mask: bool = False,
                 mask_pad_value: int = 0,
                 bgr_to_rgb: bool = False,
                 rgb_to_bgr: bool = False,
                 batch_augments: Optional[List[dict]] = None):
        super().__init__()
        assert not (bgr_to_rgb and rgb_to_bgr), (
            '`bgr2rgb` and `rgb2bgr` cannot be set to True at the same time')
        assert (mean is None) == (std is None), (
            'mean and std should be both None or tuple')
        if mean is not None:
            assert len(mean) == 3 or len(mean) == 1, (
                '`mean` should have 1 or 3 values, to be compatible with '
                f'RGB or gray image, but got {len(mean)} values')
            assert len(std) == 3 or len(std) == 1, (  # type: ignore
                '`std` should have 1 or 3 values, to be compatible with RGB '  # type: ignore # noqa: E501
                f'or gray image, but got {len(std)} values')  # type: ignore
            self._enable_normalize = True
            self.mean = np.array(mean).reshape(-1, 1, 1)
            self.std = np.array(std).reshape(-1, 1, 1)
        else:
            self._enable_normalize = False
        self._channel_conversion = rgb_to_bgr or bgr_to_rgb
        self.pad_size_divisor = pad_size_divisor
        self.pad_value = pad_value
        if batch_augments is not None:
            self.batch_augments = nn.ModuleList(
                [MODELS.build(aug) for aug in batch_augments])
        else:
            self.batch_augments = None
        self.pad_mask = pad_mask
        self.mask_pad_value = mask_pad_value

    def execute(self, data: dict, training: bool = False) -> dict:
        """Perform normalization、padding and bgr2rgb conversion based on
        ``BaseDataPreprocessor``.

        Args:
            data (dict): Data sampled from dataloader.
            training (bool): Whether to enable training time augmentation.

        Returns:
            dict: Data in the same format as the model input.
        """
        batch_pad_shape = self._get_pad_shape(data)
        _batch_inputs = data['inputs']
        # Process data with `pseudo_collate`.
        if is_list_of(_batch_inputs, jt.Var):
            batch_inputs = []
            for _batch_input in _batch_inputs:
                # channel transform
                if self._channel_conversion:
                    _batch_input = _batch_input[[2, 1, 0], ...]
                # Convert to float after channel conversion to ensure
                # efficiency
                _batch_input = _batch_input.float()
                # Normalization.
                if self._enable_normalize:
                    if self.mean.shape[0] == 3:
                        assert (
                            _batch_input.ndim == 3
                            and _batch_input.shape[0] == 3
                        ), ('If the mean has 3 values, the input tensor '
                            'should in shape of (3, H, W), but got the tensor '
                            f'with shape {_batch_input.shape}')
                    mean, std = jt.array(self.mean), jt.array(self.std)
                    _batch_input = (_batch_input - mean) / std
                batch_inputs.append(_batch_input)
            # Pad and stack Tensor.
            batch_inputs = self.stack_batch(batch_inputs,
                                            self.pad_size_divisor,
                                            self.pad_value)
        # Process data with `default_collate`.
        elif isinstance(_batch_inputs, jt.Var):
            assert _batch_inputs.ndim == 4, (
                'The input of `ImgDataPreprocessor` should be a NCHW tensor '
                'or a list of tensor, but got a tensor with shape: '
                f'{_batch_inputs.shape}')
            if self._channel_conversion:
                _batch_inputs = _batch_inputs[:, [2, 1, 0], ...]
            # Convert to float after channel conversion to ensure
            # efficiency
            _batch_inputs = _batch_inputs.float()
            if self._enable_normalize:
                mean, std = jt.array(self.mean), jt.array(self.std)
                _batch_inputs = (_batch_inputs - mean) / std
            h, w = _batch_inputs.shape[2:]
            target_h = math.ceil(
                h / self.pad_size_divisor) * self.pad_size_divisor
            target_w = math.ceil(
                w / self.pad_size_divisor) * self.pad_size_divisor
            pad_h = target_h - h
            pad_w = target_w - w
            batch_inputs = nn.pad(_batch_inputs, (0, pad_w, 0, pad_h),
                                  'constant', self.pad_value)
        else:
            raise TypeError('Output of `cast_data` should be a list of dict '
                            'or a tuple with inputs and data_samples, but got'
                            f'{type(data)}: {data}')

        data['inputs'] = batch_inputs
        data.setdefault('data_samples', None)

        inputs, data_samples = data['inputs'], data['data_samples']
        if data_samples is not None:
            # NOTE the batched image size information may be useful, e.g.
            # in DETR, this is needed for the construction of masks, which is
            # then used for the transformer_head.
            batch_input_shape = tuple(inputs[0].size()[-2:])
            for data_sample, pad_shape in zip(data_samples, batch_pad_shape):
                data_sample.set_metainfo({
                    'batch_input_shape': batch_input_shape,
                    'pad_shape': pad_shape
                })

            if self.pad_mask and training:
                self.pad_gt_masks(data_samples)

        if training and self.batch_augments is not None:
            for batch_aug in self.batch_augments:
                inputs, data_samples = batch_aug(inputs, data_samples)

        return {'inputs': inputs, 'data_samples': data_samples}

    @staticmethod
    def stack_batch(tensor_list: List[jt.Var],
                    pad_size_divisor: int = 1,
                    pad_value: Union[int, float] = 0) -> jt.Var:
        """Stack multiple tensors to form a batch and pad the tensor to the max
        shape use the right bottom padding mode in these images. If
        ``pad_size_divisor > 0``, add padding to ensure the shape of each dim
        is divisible by ``pad_size_divisor``.

        Args:
            tensor_list (List[Tensor]): A list of tensors with the same dim.
            pad_size_divisor (int): If ``pad_size_divisor > 0``, add padding
                to ensure the shape of each dim is divisible by
                ``pad_size_divisor``. This depends on the model, and many
                models need to be divisible by 32. Defaults to 1
            pad_value (int, float): The padding value. Defaults to 0.

        Returns:
        Tensor: The n dim tensor.
        """
        assert isinstance(tensor_list, list), (
            f'Expected input type to be list, but got {type(tensor_list)}')
        assert tensor_list, '`tensor_list` could not be an empty list'
        assert len({
            tensor.ndim
            for tensor in tensor_list
        }) == 1, (f'Expected the dimensions of all tensors must be the same, '
                  f'but got {[tensor.ndim for tensor in tensor_list]}')

        dim = tensor_list[0].ndim
        num_img = len(tensor_list)
        all_sizes = jt.array([tensor.shape for tensor in tensor_list])
        max_sizes = jt.ceil(
            jt.max(all_sizes, dim=0) / pad_size_divisor) * pad_size_divisor
        padded_sizes = max_sizes - all_sizes
        # The first dim normally means channel,  which should not be padded.
        padded_sizes[:, 0] = 0
        if padded_sizes.sum() == 0:
            return jt.stack(tensor_list)
        # `pad` is the second arguments of `F.pad`. If pad is (1, 2, 3, 4),
        # it means that padding the last dim with 1(left) 2(right), padding the
        # penultimate dim to 3(top) 4(bottom). The order of `pad` is opposite
        # of the `padded_sizes`. Therefore, the `padded_sizes` needs to be
        # reversed, and only odd index of pad should be assigned to keep
        # padding "right" and "bottom".
        pad = jt.zeros(num_img, 2 * dim, dtype=jt.int)
        index = list(range(dim - 1, -1, -1))
        pad[:, 1::2] = padded_sizes[:, index]
        batch_tensor = []
        for idx, tensor in enumerate(tensor_list):
            batch_tensor.append(
                nn.pad(tensor, tuple(pad[idx].tolist()), value=pad_value))
        return jt.stack(batch_tensor)

    def _get_pad_shape(self, data: dict) -> List[tuple]:
        """Get the pad_shape of each image based on data and
        pad_size_divisor."""
        _batch_inputs = data['inputs']
        # Process data with `pseudo_collate`.
        if is_list_of(_batch_inputs, jt.Var):
            batch_pad_shape = []
            for ori_input in _batch_inputs:
                pad_h = int(
                    np.ceil(ori_input.shape[1] /
                            self.pad_size_divisor)) * self.pad_size_divisor
                pad_w = int(
                    np.ceil(ori_input.shape[2] /
                            self.pad_size_divisor)) * self.pad_size_divisor
                batch_pad_shape.append((pad_h, pad_w))
        # Process data with `default_collate`.
        elif isinstance(_batch_inputs, jt.Var):
            assert _batch_inputs.ndim == 4, (
                'The input of `ImgDataPreprocessor` should be a NCHW tensor '
                'or a list of tensor, but got a tensor with shape: '
                f'{_batch_inputs.shape}')
            pad_h = int(
                np.ceil(_batch_inputs.shape[1] /
                        self.pad_size_divisor)) * self.pad_size_divisor
            pad_w = int(
                np.ceil(_batch_inputs.shape[2] /
                        self.pad_size_divisor)) * self.pad_size_divisor
            batch_pad_shape = [(pad_h, pad_w)] * _batch_inputs.shape[0]
        else:
            raise TypeError('Output of `cast_data` should be a list of dict '
                            'or a tuple with inputs and data_samples, but got'
                            f'{type(data)}: {data}')
        return batch_pad_shape

    def pad_gt_masks(self,
                     batch_data_samples: Sequence[DetDataSample]) -> None:
        """Pad gt_masks to shape of batch_input_shape."""
        if 'masks' in batch_data_samples[0].gt_instances:
            for data_samples in batch_data_samples:
                masks = data_samples.gt_instances.masks
                data_samples.gt_instances.masks = masks.pad(
                    data_samples.batch_input_shape,
                    pad_val=self.mask_pad_value)

jittordet/models/preprocessors/preprocessor.py
JittorCV/jittordet
from .base_roi_head import BaseRoIHead
from .bbox_heads import BBoxHead
from .roi_extractors import BaseRoIExtractor, SingleRoIExtractor
from .standard_roi_head import StandardRoIHead

__all__ = [
    'BaseRoIHead', 'StandardRoIHead', 'BBoxHead', 'BaseRoIExtractor',
    'SingleRoIExtractor'
]

jittordet/models/roi_heads/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/roi_heads/base_roi_head.py
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod
from typing import Tuple

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS, OptConfigType, OptMultiConfig
from jittordet.structures import InstanceList, SampleList


class BaseRoIHead(nn.Module, metaclass=ABCMeta):
    """Base class for RoIHeads."""

    def __init__(self,
                 bbox_roi_extractor: OptMultiConfig = None,
                 bbox_head: OptConfigType = None,
                 shared_head: OptConfigType = None,
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None) -> None:
        super().__init__()
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        if shared_head is not None:
            self.shared_head = MODELS.build(shared_head)

        if bbox_head is not None:
            self.init_bbox_head(bbox_roi_extractor, bbox_head)

        self.init_assigner_sampler()

    @property
    def with_bbox(self) -> bool:
        """bool: whether the RoI head contains a `bbox_head`"""
        return hasattr(self, 'bbox_head') and self.bbox_head is not None

    @property
    def with_shared_head(self) -> bool:
        """bool: whether the RoI head contains a `shared_head`"""
        return hasattr(self, 'shared_head') and self.shared_head is not None

    @abstractmethod
    def init_bbox_head(self, *args, **kwargs):
        """Initialize ``bbox_head``"""
        pass

    @abstractmethod
    def init_assigner_sampler(self, *args, **kwargs):
        """Initialize assigner and sampler."""
        pass

    @abstractmethod
    def loss(self, x: Tuple[jt.Var], rpn_results_list: InstanceList,
             batch_data_samples: SampleList):
        """Perform forward propagation and loss calculation of the roi head on
        the features of the upstream network."""

    def predict(self,
                x: Tuple[jt.Var],
                rpn_results_list: InstanceList,
                batch_data_samples: SampleList,
                rescale: bool = False) -> InstanceList:
        """Perform forward propagation of the roi head and predict detection
        results on the features of the upstream network.
        Args:
            x (tuple[Tensor]): Features from upstream network. Each
                has shape (N, C, H, W).
            rpn_results_list (list[:obj:`InstanceData`]): list of region
                proposals.
            batch_data_samples (List[:obj:`DetDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.
            rescale (bool): Whether to rescale the results to
                the original image. Defaults to True.
        Returns:
            list[obj:`InstanceData`]: Detection results of each image.
            Each item usually contains following keys.
                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
                - masks (Tensor): Has a shape (num_instances, H, W).
        """
        assert self.with_bbox, 'Bbox head must be implemented.'
        batch_img_metas = [
            data_samples.metainfo for data_samples in batch_data_samples
        ]

        # TODO: nms_op in mmcv need be enhanced, the bbox result may get
        #  difference when not rescale in bbox_head

        # If it has the mask branch, the bbox branch does not need
        # to be scaled to the original image scale, because the mask
        # branch will scale both bbox and mask at the same time.
        bbox_rescale = rescale
        results_list = self.predict_bbox(
            x,
            batch_img_metas,
            rpn_results_list,
            rcnn_test_cfg=self.test_cfg,
            rescale=bbox_rescale)
        return results_list

jittordet/models/roi_heads/base_roi_head.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple, Union

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS, TASK_UTILS, ConfigDict, ConfigType
from jittordet.models.losses import accuracy
from jittordet.models.task_utils.samplers.sampling_result import SamplingResult
from jittordet.models.utils import (empty_instances, multi_apply,
                                    multiclass_nms, normal_init)
from jittordet.structures import InstanceData, InstanceList


@MODELS.register_module()
class BBoxHead(nn.Module):
    """Simplest RoI head, with only two fc layers for classification and
    regression respectively."""

    def __init__(
        self,
        with_avg_pool: bool = False,
        with_cls: bool = True,
        with_reg: bool = True,
        roi_feat_size: int = 7,
        in_channels: int = 256,
        num_classes: int = 80,
        bbox_coder: ConfigType = dict(
            type='DeltaXYWHBBoxCoder',
            clip_border=True,
            target_means=[0., 0., 0., 0.],
            target_stds=[0.1, 0.1, 0.2, 0.2]),
        reg_class_agnostic: bool = False,
        reg_decoded_bbox: bool = False,
        reg_predictor_cfg: ConfigType = dict(type='Linear'),
        cls_predictor_cfg: ConfigType = dict(type='Linear'),
        loss_cls: ConfigType = dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
        loss_bbox: ConfigType = dict(
            type='SmoothL1Loss', beta=1.0, loss_weight=1.0)
    ) -> None:
        super().__init__()
        assert with_cls or with_reg
        self.with_avg_pool = with_avg_pool
        self.with_cls = with_cls
        self.with_reg = with_reg
        self.roi_feat_size = nn._pair(roi_feat_size)
        self.roi_feat_area = self.roi_feat_size[0] * self.roi_feat_size[1]
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.reg_class_agnostic = reg_class_agnostic
        self.reg_decoded_bbox = reg_decoded_bbox
        self.reg_predictor_cfg = reg_predictor_cfg
        self.cls_predictor_cfg = cls_predictor_cfg

        self.bbox_coder = TASK_UTILS.build(bbox_coder)
        self.loss_cls = MODELS.build(loss_cls)
        self.loss_bbox = MODELS.build(loss_bbox)

        in_channels = self.in_channels
        if self.with_avg_pool:
            self.avg_pool = nn.AvgPool2d(self.roi_feat_size)
        else:
            in_channels *= self.roi_feat_area
        if self.with_cls:
            # need to add background class
            if self.custom_cls_channels:
                cls_channels = self.loss_cls.get_cls_channels(self.num_classes)
            else:
                cls_channels = num_classes + 1
            cls_predictor_cfg_ = self.cls_predictor_cfg.copy()
            cls_predictor_cfg_.update(
                in_features=in_channels, out_features=cls_channels)
            self.fc_cls = MODELS.build(cls_predictor_cfg_)
        if self.with_reg:
            out_dim_reg = 4 if reg_class_agnostic else 4 * num_classes
            reg_predictor_cfg_ = self.reg_predictor_cfg.copy()
            reg_predictor_cfg_.update(
                in_features=in_channels, out_features=out_dim_reg)
            self.fc_reg = MODELS.build(reg_predictor_cfg_)
        self.debug_imgs = None

    def init_weights(self):
        for name, m in self.named_parameters():
            if name == 'fc_cls':
                normal_init(m, std=0.01)
            if name == 'fc_reg':
                normal_init(m, std=0.001)

    # TODO: Create a SeasawBBoxHead to simplified logic in BBoxHead
    @property
    def custom_cls_channels(self) -> bool:
        """get custom_cls_channels from loss_cls."""
        return getattr(self.loss_cls, 'custom_cls_channels', False)

    # TODO: Create a SeasawBBoxHead to simplified logic in BBoxHead
    @property
    def custom_activation(self) -> bool:
        """get custom_activation from loss_cls."""
        return getattr(self.loss_cls, 'custom_activation', False)

    # TODO: Create a SeasawBBoxHead to simplified logic in BBoxHead
    @property
    def custom_accuracy(self) -> bool:
        """get custom_accuracy from loss_cls."""
        return getattr(self.loss_cls, 'custom_accuracy', False)

    def execute(self, x: Tuple[jt.Var]) -> tuple:
        """Forward features from the upstream network.
        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.
        Returns:
            tuple: A tuple of classification scores and bbox prediction.
                - cls_score (Tensor): Classification scores for all
                  scale levels, each is a 4D-tensor, the channels number
                  is num_base_priors * num_classes.
                - bbox_pred (Tensor): Box energies / deltas for all
                  scale levels, each is a 4D-tensor, the channels number
                  is num_base_priors * 4.
        """
        if self.with_avg_pool:
            if x.numel() > 0:
                x = self.avg_pool(x)
                x = x.view(x.size(0), -1)
            else:
                # avg_pool does not support empty tensor,
                # so use torch.mean instead it
                x = jt.mean(x, dim=(-1, -2))
        cls_score = self.fc_cls(x) if self.with_cls else None
        bbox_pred = self.fc_reg(x) if self.with_reg else None
        return cls_score, bbox_pred

    def _get_targets_single(self, pos_priors: jt.Var, neg_priors: jt.Var,
                            pos_gt_bboxes: jt.Var, pos_gt_labels: jt.Var,
                            cfg: ConfigDict) -> tuple:
        """Calculate the ground truth for proposals in the single image
        according to the sampling results.
        Args:
            pos_priors (Tensor): Contains all the positive boxes,
                has shape (num_pos, 4), the last dimension 4
                represents [tl_x, tl_y, br_x, br_y].
            neg_priors (Tensor): Contains all the negative boxes,
                has shape (num_neg, 4), the last dimension 4
                represents [tl_x, tl_y, br_x, br_y].
            pos_gt_bboxes (Tensor): Contains gt_boxes for
                all positive samples, has shape (num_pos, 4),
                the last dimension 4
                represents [tl_x, tl_y, br_x, br_y].
            pos_gt_labels (Tensor): Contains gt_labels for
                all positive samples, has shape (num_pos, ).
            cfg (obj:`ConfigDict`): `train_cfg` of R-CNN.
        Returns:
            Tuple[Tensor]: Ground truth for proposals
            in a single image. Containing the following Tensors:
                - labels(Tensor): Gt_labels for all proposals, has
                  shape (num_proposals,).
                - label_weights(Tensor): Labels_weights for all
                  proposals, has shape (num_proposals,).
                - bbox_targets(Tensor):Regression target for all
                  proposals, has shape (num_proposals, 4), the
                  last dimension 4 represents [tl_x, tl_y, br_x, br_y].
                - bbox_weights(Tensor):Regression weights for all
                  proposals, has shape (num_proposals, 4).
        """
        num_pos = pos_priors.size(0)
        num_neg = neg_priors.size(0)
        num_samples = num_pos + num_neg

        # original implementation uses new_zeros since BG are set to be 0
        # now use empty & fill because BG cat_id = num_classes,
        # FG cat_id = [0, num_classes-1]
        labels = jt.full((num_samples, ), self.num_classes, dtype=jt.int64)
        label_weights = jt.zeros(num_samples)
        bbox_targets = jt.zeros(num_samples, 4)
        bbox_weights = jt.zeros(num_samples, 4)
        if num_pos > 0:
            labels[:num_pos] = pos_gt_labels
            pos_weight = 1.0 if cfg.pos_weight <= 0 else cfg.pos_weight
            label_weights[:num_pos] = pos_weight
            if not self.reg_decoded_bbox:
                pos_bbox_targets = self.bbox_coder.encode(
                    pos_priors, pos_gt_bboxes)
            else:
                # When the regression loss (e.g. `IouLoss`, `GIouLoss`)
                # is applied directly on the decoded bounding boxes, both
                # the predicted boxes and regression targets should be with
                # absolute coordinate format.
                pos_bbox_targets = pos_gt_bboxes
            bbox_targets[:num_pos, :] = pos_bbox_targets
            bbox_weights[:num_pos, :] = 1
        if num_neg > 0:
            label_weights[-num_neg:] = 1.0

        return labels, label_weights, bbox_targets, bbox_weights

    def get_targets(self,
                    sampling_results: List[SamplingResult],
                    rcnn_train_cfg: ConfigDict,
                    concat: bool = True) -> tuple:
        """Calculate the ground truth for all samples in a batch according to
        the sampling_results.
        Almost the same as the implementation in bbox_head, we passed
        additional parameters pos_inds_list and neg_inds_list to
        `_get_targets_single` function.
        Args:
            sampling_results (List[obj:SamplingResult]): Assign results of
                all images in a batch after sampling.
            rcnn_train_cfg (obj:ConfigDict): `train_cfg` of RCNN.
            concat (bool): Whether to concatenate the results of all
                the images in a single batch.
        Returns:
            Tuple[Tensor]: Ground truth for proposals in a single image.
            Containing the following list of Tensors:
            - labels (list[Tensor],Tensor): Gt_labels for all
                proposals in a batch, each tensor in list has
                shape (num_proposals,) when `concat=False`, otherwise
                just a single tensor has shape (num_all_proposals,).
            - label_weights (list[Tensor]): Labels_weights for
                all proposals in a batch, each tensor in list has
                shape (num_proposals,) when `concat=False`, otherwise
                just a single tensor has shape (num_all_proposals,).
            - bbox_targets (list[Tensor],Tensor): Regression target
                for all proposals in a batch, each tensor in list
                has shape (num_proposals, 4) when `concat=False`,
                otherwise just a single tensor has shape
                (num_all_proposals, 4), the last dimension 4 represents
                [tl_x, tl_y, br_x, br_y].
            - bbox_weights (list[tensor],Tensor): Regression weights for
                all proposals in a batch, each tensor in list has shape
                (num_proposals, 4) when `concat=False`, otherwise just a
                single tensor has shape (num_all_proposals, 4).
        """
        pos_priors_list = [res.pos_priors for res in sampling_results]
        neg_priors_list = [res.neg_priors for res in sampling_results]
        pos_gt_bboxes_list = [res.pos_gt_bboxes for res in sampling_results]
        pos_gt_labels_list = [res.pos_gt_labels for res in sampling_results]
        labels, label_weights, bbox_targets, bbox_weights = multi_apply(
            self._get_targets_single,
            pos_priors_list,
            neg_priors_list,
            pos_gt_bboxes_list,
            pos_gt_labels_list,
            cfg=rcnn_train_cfg)

        if concat:
            labels = jt.concat(labels, 0)
            label_weights = jt.concat(label_weights, 0)
            bbox_targets = jt.concat(bbox_targets, 0)
            bbox_weights = jt.concat(bbox_weights, 0)
        return labels, label_weights, bbox_targets, bbox_weights

    def loss_and_target(self,
                        cls_score: jt.Var,
                        bbox_pred: jt.Var,
                        rois: jt.Var,
                        sampling_results: List[SamplingResult],
                        rcnn_train_cfg: ConfigDict,
                        concat: bool = True,
                        reduction_override: Optional[str] = None) -> dict:
        """Calculate the loss based on the features extracted by the bbox head.
        Args:
            cls_score (Tensor): Classification prediction
                results of all class, has shape
                (batch_size * num_proposals_single_image, num_classes)
            bbox_pred (Tensor): Regression prediction results,
                has shape
                (batch_size * num_proposals_single_image, 4), the last
                dimension 4 represents [tl_x, tl_y, br_x, br_y].
            rois (Tensor): RoIs with the shape
                (batch_size * num_proposals_single_image, 5) where the first
                column indicates batch id of each RoI.
            sampling_results (List[obj:SamplingResult]): Assign results of
                all images in a batch after sampling.
            rcnn_train_cfg (obj:ConfigDict): `train_cfg` of RCNN.
            concat (bool): Whether to concatenate the results of all
                the images in a single batch. Defaults to True.
            reduction_override (str, optional): The reduction
                method used to override the original reduction
                method of the loss. Options are "none",
                "mean" and "sum". Defaults to None,
        Returns:
            dict: A dictionary of loss and targets components.
                The targets are only used for cascade rcnn.
        """

        cls_reg_targets = self.get_targets(
            sampling_results, rcnn_train_cfg, concat=concat)
        losses = self.loss(
            cls_score,
            bbox_pred,
            rois,
            *cls_reg_targets,
            reduction_override=reduction_override)

        # cls_reg_targets is only for cascade rcnn
        return dict(loss_bbox=losses, bbox_targets=cls_reg_targets)

    def loss(self,
             cls_score: jt.Var,
             bbox_pred: jt.Var,
             rois: jt.Var,
             labels: jt.Var,
             label_weights: jt.Var,
             bbox_targets: jt.Var,
             bbox_weights: jt.Var,
             reduction_override: Optional[str] = None) -> dict:
        """Calculate the loss based on the network predictions and targets.
        Args:
            cls_score (Tensor): Classification prediction
                results of all class, has shape
                (batch_size * num_proposals_single_image, num_classes)
            bbox_pred (Tensor): Regression prediction results,
                has shape
                (batch_size * num_proposals_single_image, 4), the last
                dimension 4 represents [tl_x, tl_y, br_x, br_y].
            rois (Tensor): RoIs with the shape
                (batch_size * num_proposals_single_image, 5) where the first
                column indicates batch id of each RoI.
            labels (Tensor): Gt_labels for all proposals in a batch, has
                shape (batch_size * num_proposals_single_image, ).
            label_weights (Tensor): Labels_weights for all proposals in a
                batch, has shape (batch_size * num_proposals_single_image, ).
            bbox_targets (Tensor): Regression target for all proposals in a
                batch, has shape (batch_size * num_proposals_single_image, 4),
                the last dimension 4 represents [tl_x, tl_y, br_x, br_y].
            bbox_weights (Tensor): Regression weights for all proposals in a
                batch, has shape (batch_size * num_proposals_single_image, 4).
            reduction_override (str, optional): The reduction
                method used to override the original reduction
                method of the loss. Options are "none",
                "mean" and "sum". Defaults to None,
        Returns:
            dict: A dictionary of loss.
        """

        losses = dict()

        if cls_score is not None:
            avg_factor = max(jt.sum(label_weights > 0).float().item(), 1.)
            if cls_score.numel() > 0:
                loss_cls_ = self.loss_cls(
                    cls_score,
                    labels,
                    label_weights,
                    avg_factor=avg_factor,
                    reduction_override=reduction_override)
                if isinstance(loss_cls_, dict):
                    losses.update(loss_cls_)
                else:
                    losses['loss_cls'] = loss_cls_
                if self.custom_activation:
                    acc_ = self.loss_cls.get_accuracy(cls_score, labels)
                    losses.update(acc_)
                else:
                    losses['acc'] = accuracy(cls_score, labels)
        if bbox_pred is not None:
            bg_class_ind = self.num_classes
            # 0~self.num_classes-1 are FG, self.num_classes is BG
            pos_inds = (labels >= 0) & (labels < bg_class_ind)
            # do not perform bounding box regression for BG anymore.
            if pos_inds.any():
                if self.reg_decoded_bbox:
                    # When the regression loss (e.g. `IouLoss`,
                    # `GIouLoss`, `DIouLoss`) is applied directly on
                    # the decoded bounding boxes, it decodes the
                    # already encoded coordinates to absolute format.
                    bbox_pred = self.bbox_coder.decode(rois[:, 1:], bbox_pred)
                if self.reg_class_agnostic:
                    pos_bbox_pred = bbox_pred.view(
                        bbox_pred.size(0), -1)[pos_inds.astype(jt.bool)]
                else:
                    pos_bbox_pred = bbox_pred.view(
                        bbox_pred.size(0), self.num_classes,
                        -1)[pos_inds.astype(jt.bool),
                            labels[pos_inds.astype(jt.bool)]]
                losses['loss_bbox'] = self.loss_bbox(
                    pos_bbox_pred,
                    bbox_targets[pos_inds.astype(jt.bool)],
                    bbox_weights[pos_inds.astype(jt.bool)],
                    avg_factor=bbox_targets.size(0),
                    reduction_override=reduction_override)
            else:
                losses['loss_bbox'] = bbox_pred[pos_inds].sum()

        return losses

    def predict_by_feat(self,
                        rois: Tuple[jt.Var],
                        cls_scores: Tuple[jt.Var],
                        bbox_preds: Tuple[jt.Var],
                        batch_img_metas: List[dict],
                        rcnn_test_cfg: Optional[ConfigDict] = None,
                        rescale: bool = False) -> InstanceList:
        """Transform a batch of output features extracted from the head into
        bbox results.
        Args:
            rois (tuple[Tensor]): Tuple of boxes to be transformed.
                Each has shape  (num_boxes, 5). last dimension 5 arrange as
                (batch_index, x1, y1, x2, y2).
            cls_scores (tuple[Tensor]): Tuple of box scores, each has shape
                (num_boxes, num_classes + 1).
            bbox_preds (tuple[Tensor]): Tuple of box energies / deltas, each
                has shape (num_boxes, num_classes * 4).
            batch_img_metas (list[dict]): List of image information.
            rcnn_test_cfg (obj:`ConfigDict`, optional): `test_cfg` of R-CNN.
                Defaults to None.
            rescale (bool): If True, return boxes in original image space.
                Defaults to False.
        Returns:
            list[:obj:`InstanceData`]: Instance segmentation
            results of each image after the post process.
            Each item usually contains following keys.
                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        assert len(cls_scores) == len(bbox_preds)
        result_list = []
        for img_id in range(len(batch_img_metas)):
            img_meta = batch_img_metas[img_id]
            results = self._predict_by_feat_single(
                roi=rois[img_id],
                cls_score=cls_scores[img_id],
                bbox_pred=bbox_preds[img_id],
                img_meta=img_meta,
                rescale=rescale,
                rcnn_test_cfg=rcnn_test_cfg)
            result_list.append(results)

        return result_list

    def _predict_by_feat_single(
            self,
            roi: jt.Var,
            cls_score: jt.Var,
            bbox_pred: jt.Var,
            img_meta: dict,
            rescale: bool = False,
            rcnn_test_cfg: Optional[ConfigDict] = None) -> InstanceData:
        """Transform a single image's features extracted from the head into
        bbox results.
        Args:
            roi (Tensor): Boxes to be transformed. Has shape (num_boxes, 5).
                last dimension 5 arrange as (batch_index, x1, y1, x2, y2).
            cls_score (Tensor): Box scores, has shape
                (num_boxes, num_classes + 1).
            bbox_pred (Tensor): Box energies / deltas.
                has shape (num_boxes, num_classes * 4).
            img_meta (dict): image information.
            rescale (bool): If True, return boxes in original image space.
                Defaults to False.
            rcnn_test_cfg (obj:`ConfigDict`): `test_cfg` of Bbox Head.
                Defaults to None
        Returns:
            :obj:`InstanceData`: Detection results of each image\
            Each item usually contains following keys.
                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        results = InstanceData()
        if roi.shape[0] == 0:
            return empty_instances([img_meta],
                                   instance_results=[results],
                                   num_classes=self.num_classes,
                                   score_per_cls=rcnn_test_cfg is None)[0]

        # some loss (Seesaw loss..) may have custom activation
        if self.custom_cls_channels:
            scores = self.loss_cls.get_activation(cls_score)
        else:
            scores = nn.softmax(
                cls_score, dim=-1) if cls_score is not None else None

        img_shape = img_meta['img_shape']
        num_rois = roi.size(0)
        # bbox_pred would be None in some detector when with_reg is False,
        # e.g. Grid R-CNN.
        if bbox_pred is not None:
            num_classes = 1 if self.reg_class_agnostic else self.num_classes
            roi = roi.repeat_interleave(num_classes, dim=0)
            bbox_pred = bbox_pred.view(-1, 4)
            bboxes = self.bbox_coder.decode(
                roi[..., 1:], bbox_pred, max_shape=img_shape)
        else:
            bboxes = roi[:, 1:].clone()
            if img_shape is not None and bboxes.size(-1) == 4:
                bboxes[:, [0, 2]].clamp_(min_v=0, max_v=img_shape[1])
                bboxes[:, [1, 3]].clamp_(min_v=0, max_v=img_shape[0])

        if rescale and bboxes.size(0) > 0:
            assert img_meta.get('scale_factor') is not None
            scale_factor = jt.array(
                img_meta['scale_factor'], dtype=bboxes.dtype).repeat((1, 2))
            bboxes = (bboxes.view(bboxes.size(0), -1, 4) / scale_factor).view(
                bboxes.size()[0], -1)

        # Get the inside tensor when `bboxes` is a box type
        bboxes = bboxes.view(num_rois, -1)

        if rcnn_test_cfg is None:
            # This means that it is aug test.
            # It needs to return the raw results without nms.
            results.bboxes = bboxes
            results.scores = scores
        else:
            det_bboxes, det_scores, det_labels = multiclass_nms(
                bboxes, scores, rcnn_test_cfg.score_thr, rcnn_test_cfg.nms,
                rcnn_test_cfg.max_per_img)
            results.bboxes = det_bboxes
            results.scores = det_scores
            results.labels = det_labels
        return results

    def refine_bboxes(self, sampling_results: Union[List[SamplingResult],
                                                    InstanceList],
                      bbox_results: dict,
                      batch_img_metas: List[dict]) -> InstanceList:
        """Refine bboxes during training.
        Args:
            sampling_results (List[:obj:`SamplingResult`] or
                List[:obj:`InstanceData`]): Sampling results.
                :obj:`SamplingResult` is the real sampling results
                calculate from bbox_head, while :obj:`InstanceData` is
                fake sampling results, e.g., in Sparse R-CNN or QueryInst, etc.
            bbox_results (dict): Usually is a dictionary with keys:
                - `cls_score` (Tensor): Classification scores.
                - `bbox_pred` (Tensor): Box energies / deltas.
                - `rois` (Tensor): RoIs with the shape (n, 5) where the first
                  column indicates batch id of each RoI.
                - `bbox_targets` (tuple):  Ground truth for proposals in a
                  single image. Containing the following list of Tensors:
                  (labels, label_weights, bbox_targets, bbox_weights)
            batch_img_metas (List[dict]): List of image information.
        Returns:
            list[:obj:`InstanceData`]: Refined bboxes of each image.
        Example:
            >>> # xdoctest: +REQUIRES(module:kwarray)
            >>> import numpy as np
            >>> from mmdet.models.task_modules.samplers.
            ... sampling_result import random_boxes
            >>> from mmdet.models.task_modules.samplers import SamplingResult
            >>> self = BBoxHead(reg_class_agnostic=True)
            >>> n_roi = 2
            >>> n_img = 4
            >>> scale = 512
            >>> rng = np.random.RandomState(0)
            ... batch_img_metas = [{'img_shape': (scale, scale)}
            >>>                     for _ in range(n_img)]
            >>> sampling_results = [SamplingResult.random(rng=10)
            ...                     for _ in range(n_img)]
            >>> # Create rois in the expected format
            >>> roi_boxes = random_boxes(n_roi, scale=scale, rng=rng)
            >>> img_ids = torch.randint(0, n_img, (n_roi,))
            >>> img_ids = img_ids.float()
            >>> rois = torch.cat([img_ids[:, None], roi_boxes], dim=1)
            >>> # Create other args
            >>> labels = torch.randint(0, 81, (scale,)).long()
            >>> bbox_preds = random_boxes(n_roi, scale=scale, rng=rng)
            >>> cls_score = torch.randn((scale, 81))
            ... # For each image, pretend random positive boxes are gts
            >>> bbox_targets = (labels, None, None, None)
            ... bbox_results = dict(rois=rois, bbox_pred=bbox_preds,
            ...                     cls_score=cls_score,
            ...                     bbox_targets=bbox_targets)
            >>> bboxes_list = self.refine_bboxes(sampling_results,
            ...                                  bbox_results,
            ...                                  batch_img_metas)
            >>> print(bboxes_list)
        """
        pos_is_gts = [res.pos_is_gt for res in sampling_results]
        # bbox_targets is a tuple
        labels = bbox_results['bbox_targets'][0]
        cls_scores = bbox_results['cls_score']
        rois = bbox_results['rois']
        bbox_preds = bbox_results['bbox_pred']
        if self.custom_activation:
            # TODO: Create a SeasawBBoxHead to simplified logic in BBoxHead
            cls_scores = self.loss_cls.get_activation(cls_scores)
        if cls_scores.numel() == 0:
            return None
        if cls_scores.shape[-1] == self.num_classes + 1:
            # remove background class
            cls_scores = cls_scores[:, :-1]
        elif cls_scores.shape[-1] != self.num_classes:
            raise ValueError('The last dim of `cls_scores` should equal to '
                             '`num_classes` or `num_classes + 1`,'
                             f'but got {cls_scores.shape[-1]}.')
        labels = jt.where(labels == self.num_classes, cls_scores.argmax(1),
                          labels)

        img_ids = rois[:, 0].long().unique(sorted=True)
        assert img_ids.numel() <= len(batch_img_metas)

        results_list = []
        for i in range(len(batch_img_metas)):
            inds = jt.nonzero(rois[:, 0] == i).squeeze(dim=1)
            num_rois = inds.numel()

            bboxes_ = rois[inds, 1:]
            label_ = labels[inds]
            bbox_pred_ = bbox_preds[inds]
            img_meta_ = batch_img_metas[i]
            pos_is_gts_ = pos_is_gts[i]

            bboxes = self.regress_by_class(bboxes_, label_, bbox_pred_,
                                           img_meta_)
            # filter gt bboxes
            pos_keep = 1 - pos_is_gts_
            keep_inds = pos_is_gts_.new_ones(num_rois)
            keep_inds[:len(pos_is_gts_)] = pos_keep
            results = InstanceData(bboxes=bboxes[keep_inds.astype(jt.bool)])
            results_list.append(results)

        return results_list

    def regress_by_class(self, priors: jt.Var, label: jt.Var,
                         bbox_pred: jt.Var, img_meta: dict) -> jt.Var:
        """Regress the bbox for the predicted class. Used in Cascade R-CNN.
        Args:
            priors (Tensor): Priors from `rpn_head` or last stage
                `bbox_head`, has shape (num_proposals, 4).
            label (Tensor): Only used when `self.reg_class_agnostic`
                is False, has shape (num_proposals, ).
            bbox_pred (Tensor): Regression prediction of
                current stage `bbox_head`. When `self.reg_class_agnostic`
                is False, it has shape (n, num_classes * 4), otherwise
                it has shape (n, 4).
            img_meta (dict): Image meta info.
        Returns:
            Tensor: Regressed bboxes, the same shape as input rois.
        """
        if not self.reg_class_agnostic:
            label = label * 4
            inds = jt.stack([label + i for i in range(4)], 1)
            bbox_pred = jt.gather(bbox_pred, 1, inds)
        assert bbox_pred.size()[1] == 4

        max_shape = img_meta['img_shape']
        regressed_bboxes = self.bbox_coder.decode(
            priors, bbox_pred, max_shape=max_shape)
        return regressed_bboxes

jittordet/models/roi_heads/bbox_heads/bbox_head.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Tuple, Union

import jittor as jt
import jittor.nn as nn

from jittordet.engine import MODELS, ConfigDict
from jittordet.models.layers import ConvModule
from jittordet.models.utils import xavier_init
from .bbox_head import BBoxHead


@MODELS.register_module()
class ConvFCBBoxHead(BBoxHead):
    r"""More general bbox head, with shared conv and fc layers and two optional
    separated branches.
    .. code-block:: none
                                    /-> cls convs -> cls fcs -> cls
        shared convs -> shared fcs
                                    \-> reg convs -> reg fcs -> reg
    """  # noqa: W605

    def __init__(self,
                 num_shared_convs: int = 0,
                 num_shared_fcs: int = 0,
                 num_cls_convs: int = 0,
                 num_cls_fcs: int = 0,
                 num_reg_convs: int = 0,
                 num_reg_fcs: int = 0,
                 conv_out_channels: int = 256,
                 fc_out_channels: int = 1024,
                 conv_cfg: Optional[Union[dict, ConfigDict]] = None,
                 norm_cfg: Optional[Union[dict, ConfigDict]] = None,
                 *args,
                 **kwargs) -> None:
        super().__init__(*args, **kwargs)
        assert (num_shared_convs + num_shared_fcs + num_cls_convs +
                num_cls_fcs + num_reg_convs + num_reg_fcs > 0)
        if num_cls_convs > 0 or num_reg_convs > 0:
            assert num_shared_fcs == 0
        if not self.with_cls:
            assert num_cls_convs == 0 and num_cls_fcs == 0
        if not self.with_reg:
            assert num_reg_convs == 0 and num_reg_fcs == 0
        self.num_shared_convs = num_shared_convs
        self.num_shared_fcs = num_shared_fcs
        self.num_cls_convs = num_cls_convs
        self.num_cls_fcs = num_cls_fcs
        self.num_reg_convs = num_reg_convs
        self.num_reg_fcs = num_reg_fcs
        self.conv_out_channels = conv_out_channels
        self.fc_out_channels = fc_out_channels
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg

        # add shared convs and fcs
        self.shared_convs, self.shared_fcs, last_layer_dim = \
            self._add_conv_fc_branch(
                self.num_shared_convs, self.num_shared_fcs, self.in_channels,
                True)
        self.shared_out_channels = last_layer_dim

        # add cls specific branch
        self.cls_convs, self.cls_fcs, self.cls_last_dim = \
            self._add_conv_fc_branch(
                self.num_cls_convs, self.num_cls_fcs, self.shared_out_channels)

        # add reg specific branch
        self.reg_convs, self.reg_fcs, self.reg_last_dim = \
            self._add_conv_fc_branch(
                self.num_reg_convs, self.num_reg_fcs, self.shared_out_channels)

        if self.num_shared_fcs == 0 and not self.with_avg_pool:
            if self.num_cls_fcs == 0:
                self.cls_last_dim *= self.roi_feat_area
            if self.num_reg_fcs == 0:
                self.reg_last_dim *= self.roi_feat_area

        self.relu = nn.ReLU()
        # reconstruct fc_cls and fc_reg since input channels are changed
        if self.with_cls:
            if self.custom_cls_channels:
                cls_channels = self.loss_cls.get_cls_channels(self.num_classes)
            else:
                cls_channels = self.num_classes + 1
            cls_predictor_cfg_ = self.cls_predictor_cfg.copy()
            cls_predictor_cfg_.update(
                in_features=self.cls_last_dim, out_features=cls_channels)
            self.fc_cls = MODELS.build(cls_predictor_cfg_)
        if self.with_reg:
            out_dim_reg = 4 if self.reg_class_agnostic else \
                4 * self.num_classes
            reg_predictor_cfg_ = self.reg_predictor_cfg.copy()
            reg_predictor_cfg_.update(
                in_features=self.reg_last_dim, out_features=out_dim_reg)
            self.fc_reg = MODELS.build(reg_predictor_cfg_)

    def init_weights(self):
        super().init_weights()
        for name, m in self.named_parameters():
            if name in ['shared_fcs', 'cls_fcs', 'reg_fcs']:
                xavier_init(m, distribution='uniform')

    def _add_conv_fc_branch(self,
                            num_branch_convs: int,
                            num_branch_fcs: int,
                            in_channels: int,
                            is_shared: bool = False) -> tuple:
        """Add shared or separable branch.

        convs -> avg pool (optional) -> fcs
        """
        last_layer_dim = in_channels
        # add branch specific conv layers
        branch_convs = nn.ModuleList()
        if num_branch_convs > 0:
            for i in range(num_branch_convs):
                conv_in_channels = (
                    last_layer_dim if i == 0 else self.conv_out_channels)
                branch_convs.append(
                    ConvModule(
                        conv_in_channels,
                        self.conv_out_channels,
                        3,
                        padding=1,
                        conv_cfg=self.conv_cfg,
                        norm_cfg=self.norm_cfg))
            last_layer_dim = self.conv_out_channels
        # add branch specific fc layers
        branch_fcs = nn.ModuleList()
        if num_branch_fcs > 0:
            # for shared branch, only consider self.with_avg_pool
            # for separated branches, also consider self.num_shared_fcs
            if (is_shared
                    or self.num_shared_fcs == 0) and not self.with_avg_pool:
                last_layer_dim *= self.roi_feat_area
            for i in range(num_branch_fcs):
                fc_in_channels = (
                    last_layer_dim if i == 0 else self.fc_out_channels)
                branch_fcs.append(
                    nn.Linear(fc_in_channels, self.fc_out_channels))
            last_layer_dim = self.fc_out_channels
        return branch_convs, branch_fcs, last_layer_dim

    def execute(self, x: Tuple[jt.Var]) -> tuple:
        """Forward features from the upstream network.
        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.
        Returns:
            tuple: A tuple of classification scores and bbox prediction.
                - cls_score (Tensor): Classification scores for all \
                    scale levels, each is a 4D-tensor, the channels number \
                    is num_base_priors * num_classes.
                - bbox_pred (Tensor): Box energies / deltas for all \
                    scale levels, each is a 4D-tensor, the channels number \
                    is num_base_priors * 4.
        """
        # shared part
        if self.num_shared_convs > 0:
            for conv in self.shared_convs:
                x = conv(x)

        if self.num_shared_fcs > 0:
            if self.with_avg_pool:
                x = self.avg_pool(x)

            x = x.flatten(1)

            for fc in self.shared_fcs:
                x = self.relu(fc(x))
        # separate branches
        x_cls = x
        x_reg = x

        for conv in self.cls_convs:
            x_cls = conv(x_cls)
        if len(x_cls.size()) > 2:
            if self.with_avg_pool:
                x_cls = self.avg_pool(x_cls)
            x_cls = x_cls.flatten(1)
        for fc in self.cls_fcs:
            x_cls = self.relu(fc(x_cls))

        for conv in self.reg_convs:
            x_reg = conv(x_reg)
        if len(x_reg.size()) > 2:
            if self.with_avg_pool:
                x_reg = self.avg_pool(x_reg)
            x_reg = x_reg.flatten(1)
        for fc in self.reg_fcs:
            x_reg = self.relu(fc(x_reg))

        cls_score = self.fc_cls(x_cls) if self.with_cls else None
        bbox_pred = self.fc_reg(x_reg) if self.with_reg else None
        return cls_score, bbox_pred


@MODELS.register_module()
class Shared2FCBBoxHead(ConvFCBBoxHead):

    def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None:
        super().__init__(
            num_shared_convs=0,
            num_shared_fcs=2,
            num_cls_convs=0,
            num_cls_fcs=0,
            num_reg_convs=0,
            num_reg_fcs=0,
            fc_out_channels=fc_out_channels,
            *args,
            **kwargs)


@MODELS.register_module()
class Shared4Conv1FCBBoxHead(ConvFCBBoxHead):

    def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None:
        super().__init__(
            num_shared_convs=4,
            num_shared_fcs=1,
            num_cls_convs=0,
            num_cls_fcs=0,
            num_reg_convs=0,
            num_reg_fcs=0,
            fc_out_channels=fc_out_channels,
            *args,
            **kwargs)

jittordet/models/roi_heads/bbox_heads/convfc_bbox_head.py
JittorCV/jittordet
from .base_roi_extractor import BaseRoIExtractor
from .single_level_roi_extractor import SingleRoIExtractor

__all__ = ['BaseRoIExtractor', 'SingleRoIExtractor']

jittordet/models/roi_heads/roi_extractors/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/roi_heads/roi_extractors/base_roi_extractor.py
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod
from typing import List, Optional, Tuple

import jittor as jt
import jittor.nn as nn

from jittordet import ops
from jittordet.engine import ConfigType


class BaseRoIExtractor(nn.Module, metaclass=ABCMeta):
    """Base class for RoI extractor.

    Args:
        roi_layer (:obj:`ConfigDict` or dict): Specify RoI layer type and
            arguments.
        out_channels (int): Output channels of RoI layers.
        featmap_strides (list[int]): Strides of input feature maps.
        init_cfg (:obj:`ConfigDict` or dict or list[:obj:`ConfigDict` or \
            dict], optional): Initialization config dict. Defaults to None.
    """

    def __init__(self, roi_layer: ConfigType, out_channels: int,
                 featmap_strides: List[int]) -> None:
        super().__init__()
        self.roi_layers = self.build_roi_layers(roi_layer, featmap_strides)
        self.out_channels = out_channels
        self.featmap_strides = featmap_strides

    @property
    def num_inputs(self) -> int:
        """int: Number of input feature maps."""
        return len(self.featmap_strides)

    def build_roi_layers(self, layer_cfg: ConfigType,
                         featmap_strides: List[int]) -> nn.ModuleList:
        """Build RoI operator to extract feature from each level feature map.

        Args:
            layer_cfg (:obj:`ConfigDict` or dict): Dictionary to construct and
                config RoI layer operation. Options are modules under
                ``mmcv/ops`` such as ``RoIAlign``.
            featmap_strides (list[int]): The stride of input feature map w.r.t
                to the original image size, which would be used to scale RoI
                coordinate (original image coordinate system) to feature
                coordinate system.
        Returns:
            :obj:`nn.ModuleList`: The RoI extractor modules for each level
                feature map.
        """

        cfg = layer_cfg.copy()
        layer_type = cfg.pop('type')
        assert hasattr(ops, layer_type)
        layer_cls = getattr(ops, layer_type)
        roi_layers = nn.ModuleList(
            [layer_cls(spatial_scale=1 / s, **cfg) for s in featmap_strides])
        return roi_layers

    def roi_rescale(self, rois: jt.Var, scale_factor: float) -> jt.Var:
        """Scale RoI coordinates by scale factor.

        Args:
            rois (Tensor): RoI (Region of Interest), shape (n, 5)
            scale_factor (float): Scale factor that RoI will be multiplied by.
        Returns:
            Tensor: Scaled RoI.
        """

        cx = (rois[:, 1] + rois[:, 3]) * 0.5
        cy = (rois[:, 2] + rois[:, 4]) * 0.5
        w = rois[:, 3] - rois[:, 1]
        h = rois[:, 4] - rois[:, 2]
        new_w = w * scale_factor
        new_h = h * scale_factor
        x1 = cx - new_w * 0.5
        x2 = cx + new_w * 0.5
        y1 = cy - new_h * 0.5
        y2 = cy + new_h * 0.5
        new_rois = jt.stack((rois[:, 0], x1, y1, x2, y2), dim=-1)
        return new_rois

    @abstractmethod
    def execute(self,
                feats: Tuple[jt.Var],
                rois: jt.Var,
                roi_scale_factor: Optional[float] = None) -> jt.Var:
        """Extractor ROI feats.

        Args:
            feats (Tuple[Tensor]): Multi-scale features.
            rois (Tensor): RoIs with the shape (n, 5) where the first
                column indicates batch id of each RoI.
            roi_scale_factor (Optional[float]): RoI scale factor.
                Defaults to None.
        Returns:
            Tensor: RoI feature.
        """
        pass

jittordet/models/roi_heads/roi_extractors/base_roi_extractor.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple

import jittor as jt

from jittordet.engine import MODELS, ConfigType
from .base_roi_extractor import BaseRoIExtractor


@MODELS.register_module()
class SingleRoIExtractor(BaseRoIExtractor):
    """Extract RoI features from a single level feature map.

    If there are multiple input feature levels, each RoI is mapped to a level
    according to its scale. The mapping rule is proposed in
    `FPN <https://arxiv.org/abs/1612.03144>`_.
    Args:
        roi_layer (:obj:`ConfigDict` or dict): Specify RoI layer type and
            arguments.
        out_channels (int): Output channels of RoI layers.
        featmap_strides (List[int]): Strides of input feature maps.
        finest_scale (int): Scale threshold of mapping to level 0.
            Defaults to 56.
        init_cfg (:obj:`ConfigDict` or dict or list[:obj:`ConfigDict` or \
            dict], optional): Initialization config dict. Defaults to None.
    """

    def __init__(self,
                 roi_layer: ConfigType,
                 out_channels: int,
                 featmap_strides: List[int],
                 finest_scale: int = 56) -> None:
        super().__init__(
            roi_layer=roi_layer,
            out_channels=out_channels,
            featmap_strides=featmap_strides)
        self.finest_scale = finest_scale

    def map_roi_levels(self, rois: jt.Var, num_levels: int) -> jt.Var:
        """Map rois to corresponding feature levels by scales.

        - scale < finest_scale * 2: level 0
        - finest_scale * 2 <= scale < finest_scale * 4: level 1
        - finest_scale * 4 <= scale < finest_scale * 8: level 2
        - scale >= finest_scale * 8: level 3
        Args:
            rois (Tensor): Input RoIs, shape (k, 5).
            num_levels (int): Total level number.
        Returns:
            Tensor: Level index (0-based) of each RoI, shape (k, )
        """
        scale = jt.sqrt((rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2]))
        target_lvls = jt.floor(jt.log2(scale / self.finest_scale + 1e-6))
        target_lvls = target_lvls.clamp(
            min_v=0, max_v=num_levels - 1).astype(jt.int64)
        return target_lvls

    def execute(self,
                feats: Tuple[jt.Var],
                rois: jt.Var,
                roi_scale_factor: Optional[float] = None):
        """Extractor ROI feats.

        Args:
            feats (Tuple[Tensor]): Multi-scale features.
            rois (Tensor): RoIs with the shape (n, 5) where the first
                column indicates batch id of each RoI.
            roi_scale_factor (Optional[float]): RoI scale factor.
                Defaults to None.
        Returns:
            Tensor: RoI feature.
        """
        # convert fp32 to fp16 when amp is on
        rois = rois.type_as(feats[0])
        out_size = self.roi_layers[0].output_size
        num_levels = len(feats)
        roi_feats = jt.zeros(
            rois.size(0), self.out_channels, *out_size, dtype=feats[0].dtype)

        if num_levels == 1:
            if len(rois) == 0:
                return roi_feats
            return self.roi_layers[0](feats[0], rois)

        target_lvls = self.map_roi_levels(rois, num_levels)

        if roi_scale_factor is not None:
            rois = self.roi_rescale(rois, roi_scale_factor)
        for i in range(num_levels):
            mask = target_lvls == i
            inds = mask.nonzero().squeeze(1)
            if inds.numel() > 0:
                rois_ = rois[inds]
                roi_feats_t = self.roi_layers[i](feats[i], rois_)
                roi_feats[inds] = roi_feats_t
            else:
                # Sometimes some pyramid levels will not be used for RoI
                # feature extraction and this will cause an incomplete
                # computation graph in one GPU, which is different from those
                # in other GPUs and will cause a hanging error.
                # Therefore, we add it to ensure each feature pyramid is
                # included in the computation graph to avoid runtime bugs.
                roi_feats += sum(
                    x.view(-1)[0]
                    for x in self.parameters()) * 0. + feats[i].sum() * 0.
        return roi_feats

jittordet/models/roi_heads/roi_extractors/single_level_roi_extractor.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/models/roi_heads/standard_roi_head.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple

import jittor as jt

from jittordet.engine import MODELS, TASK_UTILS, ConfigType
from jittordet.structures import DetDataSample, InstanceList
from ..task_utils.samplers.sampling_result import SamplingResult
from ..utils import bbox2roi, empty_instances, unpack_gt_instances
from .base_roi_head import BaseRoIHead


@MODELS.register_module()
class StandardRoIHead(BaseRoIHead):
    """Simplest base roi head including one bbox head."""

    def init_assigner_sampler(self) -> None:
        """Initialize assigner and sampler."""
        self.bbox_assigner = None
        self.bbox_sampler = None
        if self.train_cfg:
            self.bbox_assigner = TASK_UTILS.build(self.train_cfg.assigner)
            self.bbox_sampler = TASK_UTILS.build(
                self.train_cfg.sampler, default_args=dict(context=self))

    def init_bbox_head(self, bbox_roi_extractor: ConfigType,
                       bbox_head: ConfigType) -> None:
        """Initialize box head and box roi extractor.

        Args:
            bbox_roi_extractor (dict or ConfigDict): Config of box
                roi extractor.
            bbox_head (dict or ConfigDict): Config of box in box head.
        """
        self.bbox_roi_extractor = MODELS.build(bbox_roi_extractor)
        self.bbox_head = MODELS.build(bbox_head)

    def execute(self, x: Tuple[jt.Var],
                rpn_results_list: InstanceList) -> tuple:
        """Network forward process.

        Usually includes backbone, neck and head
        forward without any post-processing.
        Args:
            x (List[Tensor]): Multi-level features that may have different
                resolutions.
            rpn_results_list (list[:obj:`InstanceData`]): List of region
                proposals.
        Returns
            tuple: A tuple of features from ``bbox_head``
            forward.
        """
        results = ()
        proposals = [rpn_results.bboxes for rpn_results in rpn_results_list]
        rois = bbox2roi(proposals)
        # bbox head
        if self.with_bbox:
            bbox_results = self._bbox_execute(x, rois)
            results = results + (bbox_results['cls_score'],
                                 bbox_results['bbox_pred'])
        return results

    def loss(self, x: Tuple[jt.Var], rpn_results_list: InstanceList,
             batch_data_samples: List[DetDataSample]) -> dict:
        """Perform forward propagation and loss calculation of the detection
        roi on the features of the upstream network.

        Args:
            x (tuple[Tensor]): List of multi-level img features.
            rpn_results_list (list[:obj:`InstanceData`]): List of region
                proposals.
            batch_data_samples (list[:obj:`DetDataSample`]): The batch
                data samples. It usually includes information such
                as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`.
        Returns:
            dict[str, Tensor]: A dictionary of loss components
        """
        assert len(rpn_results_list) == len(batch_data_samples)
        outputs = unpack_gt_instances(batch_data_samples)
        (batch_gt_instances, batch_img_metas,
         batch_gt_instances_ignore) = outputs

        # assign gts and sample proposals
        num_imgs = len(batch_data_samples)
        sampling_results = []
        for i in range(num_imgs):
            # rename rpn_results.bboxes to rpn_results.priors
            rpn_results = rpn_results_list[i]
            rpn_results.priors = rpn_results.pop('bboxes')

            assign_result = self.bbox_assigner.assign(
                rpn_results, batch_gt_instances[i],
                batch_gt_instances_ignore[i])
            sampling_result = self.bbox_sampler.sample(
                assign_result,
                rpn_results,
                batch_gt_instances[i],
                feats=[lvl_feat[i][None] for lvl_feat in x])
            sampling_results.append(sampling_result)

        losses = dict()
        # bbox head loss
        if self.with_bbox:
            bbox_results = self.bbox_loss(x, sampling_results)
            losses.update(bbox_results['loss_bbox'])
        return losses

    def _bbox_execute(self, x: Tuple[jt.Var], rois: jt.Var) -> dict:
        """Box head forward function used in both training and testing.
        Args:
            x (tuple[Tensor]): List of multi-level img features.
            rois (Tensor): RoIs with the shape (n, 5) where the first
                column indicates batch id of each RoI.
        Returns:
             dict[str, Tensor]: Usually returns a dictionary with keys:
                - `cls_score` (Tensor): Classification scores.
                - `bbox_pred` (Tensor): Box energies / deltas.
                - `bbox_feats` (Tensor): Extract bbox RoI features.
        """
        # TODO: a more flexible way to decide which feature maps to use
        bbox_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs], rois)
        if self.with_shared_head:
            bbox_feats = self.shared_head(bbox_feats)
        cls_score, bbox_pred = self.bbox_head(bbox_feats)

        bbox_results = dict(
            cls_score=cls_score, bbox_pred=bbox_pred, bbox_feats=bbox_feats)
        return bbox_results

    def bbox_loss(self, x: Tuple[jt.Var],
                  sampling_results: List[SamplingResult]) -> dict:
        """Perform forward propagation and loss calculation of the bbox head on
        the features of the upstream network.
        Args:
            x (tuple[Tensor]): List of multi-level img features.
            sampling_results (list["obj:`SamplingResult`]): Sampling results.
        Returns:
            dict[str, Tensor]: Usually returns a dictionary with keys:
                - `cls_score` (Tensor): Classification scores.
                - `bbox_pred` (Tensor): Box energies / deltas.
                - `bbox_feats` (Tensor): Extract bbox RoI features.
                - `loss_bbox` (dict): A dictionary of bbox loss components.
        """
        rois = bbox2roi([res.priors for res in sampling_results])
        bbox_results = self._bbox_execute(x, rois)

        bbox_loss_and_target = self.bbox_head.loss_and_target(
            cls_score=bbox_results['cls_score'],
            bbox_pred=bbox_results['bbox_pred'],
            rois=rois,
            sampling_results=sampling_results,
            rcnn_train_cfg=self.train_cfg)
        bbox_results.update(loss_bbox=bbox_loss_and_target['loss_bbox'])
        return bbox_results

    def predict_bbox(self,
                     x: Tuple[jt.Var],
                     batch_img_metas: List[dict],
                     rpn_results_list: InstanceList,
                     rcnn_test_cfg: ConfigType,
                     rescale: bool = False) -> InstanceList:
        """Perform forward propagation of the bbox head and predict detection
        results on the features of the upstream network.
        Args:
            x (tuple[Tensor]): Feature maps of all scale level.
            batch_img_metas (list[dict]): List of image information.
            rpn_results_list (list[:obj:`InstanceData`]): List of region
                proposals.
            rcnn_test_cfg (obj:`ConfigDict`): `test_cfg` of R-CNN.
            rescale (bool): If True, return boxes in original image space.
                Defaults to False.
        Returns:
            list[:obj:`InstanceData`]: Detection results of each image
            after the post process.
            Each item usually contains following keys.
                - scores (Tensor): Classification scores, has a shape
                  (num_instance, )
                - labels (Tensor): Labels of bboxes, has a shape
                  (num_instances, ).
                - bboxes (Tensor): Has a shape (num_instances, 4),
                  the last dimension 4 arrange as (x1, y1, x2, y2).
        """
        proposals = [res.bboxes for res in rpn_results_list]
        rois = bbox2roi(proposals)

        if rois.shape[0] == 0:
            return empty_instances(
                batch_img_metas,
                num_classes=self.bbox_head.num_classes,
                score_per_cls=rcnn_test_cfg is None)

        bbox_results = self._bbox_execute(x, rois)

        # split batch bbox prediction back to each image
        cls_scores = bbox_results['cls_score']
        bbox_preds = bbox_results['bbox_pred']
        num_proposals_per_img = tuple(len(p) for p in proposals)
        rois = rois.split(num_proposals_per_img, 0)
        cls_scores = cls_scores.split(num_proposals_per_img, 0)

        # some detector with_reg is False, bbox_preds will be None
        if bbox_preds is not None:
            # TODO move this to a sabl_roi_head
            # the bbox prediction of some detectors like SABL is not Tensor
            if isinstance(bbox_preds, jt.Var):
                bbox_preds = bbox_preds.split(num_proposals_per_img, 0)
            else:
                bbox_preds = self.bbox_head.bbox_pred_split(
                    bbox_preds, num_proposals_per_img)
        else:
            bbox_preds = (None, ) * len(proposals)

        result_list = self.bbox_head.predict_by_feat(
            rois=rois,
            cls_scores=cls_scores,
            bbox_preds=bbox_preds,
            batch_img_metas=batch_img_metas,
            rcnn_test_cfg=rcnn_test_cfg,
            rescale=rescale)
        return result_list

jittordet/models/roi_heads/standard_roi_head.py
JittorCV/jittordet
from .assigners import (AssignResult, BaseAssigner, BboxOverlaps2D,
                        MaxIoUAssigner)
from .bbox_coders import BaseBBoxCoder, DeltaXYWHBBoxCoder
from .prior_generators import AnchorGenerator, anchor_inside_flags
from .samplers import BaseSampler, PseudoSampler, RandomSampler, SamplingResult

__all__ = [
    'AnchorGenerator', 'anchor_inside_flags', 'BaseAssigner', 'MaxIoUAssigner',
    'AssignResult', 'BboxOverlaps2D', 'BaseBBoxCoder', 'DeltaXYWHBBoxCoder',
    'BaseSampler', 'PseudoSampler', 'RandomSampler', 'SamplingResult'
]

jittordet/models/task_utils/__init__.py
JittorCV/jittordet
from .assign_result import AssignResult
from .base_assigner import BaseAssigner
from .iou2d_calculator import BboxOverlaps2D
from .max_iou_assigner import MaxIoUAssigner

__all__ = ['BaseAssigner', 'MaxIoUAssigner', 'AssignResult', 'BboxOverlaps2D']

jittordet/models/task_utils/assigners/__init__.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt


class AssignResult:
    """Stores assignments between predicted and truth boxes.

    Attributes:
        num_gts (int): the number of truth boxes considered when computing this
            assignment
        gt_inds (Tensor): for each predicted box indicates the 1-based
            index of the assigned truth box. 0 means unassigned and -1 means
            ignore.
        max_overlaps (Tensor): the iou between the predicted box and its
            assigned truth box.
        labels (Tensor): If specified, for each predicted box
            indicates the category label of the assigned truth box.

    Example:
        >>> # An assign result between 4 predicted boxes and 9 true boxes
        >>> # where only two boxes were assigned.
        >>> num_gts = 9
        >>> max_overlaps = torch.LongTensor([0, .5, .9, 0])
        >>> gt_inds = torch.LongTensor([-1, 1, 2, 0])
        >>> labels = torch.LongTensor([0, 3, 4, 0])
        >>> self = AssignResult(num_gts, gt_inds, max_overlaps, labels)
        >>> print(str(self))  # xdoctest: +IGNORE_WANT
        <AssignResult(num_gts=9, gt_inds.shape=(4,), max_overlaps.shape=(4,),
                      labels.shape=(4,))>
        >>> # Force addition of gt labels (when adding gt as proposals)
        >>> new_labels = torch.LongTensor([3, 4, 5])
        >>> self.add_gt_(new_labels)
        >>> print(str(self))  # xdoctest: +IGNORE_WANT
        <AssignResult(num_gts=9, gt_inds.shape=(7,), max_overlaps.shape=(7,),
                      labels.shape=(7,))>
    """

    def __init__(self, num_gts: int, gt_inds: jt.Var, max_overlaps: jt.Var,
                 labels: jt.Var) -> None:
        self.num_gts = num_gts
        self.gt_inds = gt_inds
        self.max_overlaps = max_overlaps
        self.labels = labels
        # Interface for possible user-defined properties
        self._extra_properties = {}

    @property
    def num_preds(self):
        """int: the number of predictions in this assignment"""
        return len(self.gt_inds)

    def set_extra_property(self, key, value):
        """Set user-defined new property."""
        assert key not in self.info
        self._extra_properties[key] = value

    def get_extra_property(self, key):
        """Get user-defined property."""
        return self._extra_properties.get(key, None)

    @property
    def info(self):
        """dict: a dictionary of info about the object"""
        basic_info = {
            'num_gts': self.num_gts,
            'num_preds': self.num_preds,
            'gt_inds': self.gt_inds,
            'max_overlaps': self.max_overlaps,
            'labels': self.labels,
        }
        basic_info.update(self._extra_properties)
        return basic_info

    def add_gt_(self, gt_labels):
        """Add ground truth as assigned results.

        Args:
            gt_labels (torch.Tensor): Labels of gt boxes
        """
        self_inds = jt.arange(1, len(gt_labels) + 1, dtype=jt.int64)
        self.gt_inds = jt.concat([self_inds, self.gt_inds])

        self.max_overlaps = jt.concat([
            jt.ones(len(gt_labels), dtype=self.max_overlaps.dtype),
            self.max_overlaps
        ])

        self.labels = jt.concat([gt_labels, self.labels])

jittordet/models/task_utils/assigners/assign_result.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod
from typing import Optional

from jittordet.structures import InstanceData


class BaseAssigner(metaclass=ABCMeta):
    """Base assigner that assigns boxes to ground truth boxes."""

    @abstractmethod
    def assign(self,
               pred_instances: InstanceData,
               gt_instances: InstanceData,
               gt_instances_ignore: Optional[InstanceData] = None,
               **kwargs):
        """Assign boxes to either a ground truth boxes or a negative boxes."""

jittordet/models/task_utils/assigners/base_assigner.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/assigners/iou2d_calculator.py
# Copyright (c) OpenMMLab. All rights reserved.
from jittordet.engine import TASK_UTILS
from jittordet.utils import bbox_overlaps


@TASK_UTILS.register_module()
class BboxOverlaps2D:
    """2D Overlaps (e.g. IoUs, GIoUs) Calculator."""

    def __call__(self, bboxes1, bboxes2, mode='iou', is_aligned=False):
        """Calculate IoU between 2D bboxes.

        Args:
            bboxes1 (Tensor or :obj:`BaseBoxes`): bboxes have shape (m, 4)
                in <x1, y1, x2, y2> format, or shape (m, 5) in <x1, y1, x2,
                y2, score> format.
            bboxes2 (Tensor or :obj:`BaseBoxes`): bboxes have shape (m, 4)
                in <x1, y1, x2, y2> format, shape (m, 5) in <x1, y1, x2, y2,
                score> format, or be empty. If ``is_aligned `` is ``True``,
                then m and n must be equal.
            mode (str): "iou" (intersection over union), "iof" (intersection
                over foreground), or "giou" (generalized intersection over
                union).
            is_aligned (bool, optional): If True, then m and n must be equal.
                Default False.

        Returns:
            Tensor: shape (m, n) if ``is_aligned `` is False else shape (m,)
        """
        assert bboxes1.size(-1) in [0, 4, 5]
        assert bboxes2.size(-1) in [0, 4, 5]
        if bboxes2.size(-1) == 5:
            bboxes2 = bboxes2[..., :4]
        if bboxes1.size(-1) == 5:
            bboxes1 = bboxes1[..., :4]

        return bbox_overlaps(bboxes1, bboxes2, mode, is_aligned)

jittordet/models/task_utils/assigners/iou2d_calculator.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/assigners/max_iou_assigner.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Union

import jittor as jt

from jittordet.engine import TASK_UTILS
from jittordet.structures import InstanceData
from .assign_result import AssignResult
from .base_assigner import BaseAssigner


@TASK_UTILS.register_module()
class MaxIoUAssigner(BaseAssigner):
    """Assign a corresponding gt bbox or background to each bbox.

    Each proposals will be assigned with `-1`, or a semi-positive integer
    indicating the ground truth index.

    - -1: negative sample, no assigned gt
    - semi-positive integer: positive sample, index (0-based) of assigned gt

    Args:
        pos_iou_thr (float): IoU threshold for positive bboxes.
        neg_iou_thr (float or tuple): IoU threshold for negative bboxes.
        min_pos_iou (float): Minimum iou for a bbox to be considered as a
            positive bbox. Positive samples can have smaller IoU than
            pos_iou_thr due to the 4th step (assign max IoU sample to each gt).
            `min_pos_iou` is set to avoid assigning bboxes that have extremely
            small iou with GT as positive samples. It brings about 0.3 mAP
            improvements in 1x schedule but does not affect the performance of
            3x schedule. More comparisons can be found in
            `PR #7464 <https://github.com/open-mmlab/mmdetection/pull/7464>`_.
        gt_max_assign_all (bool): Whether to assign all bboxes with the same
            highest overlap with some gt to that gt.
        ignore_iof_thr (float): IoF threshold for ignoring bboxes (if
            `gt_bboxes_ignore` is specified). Negative values mean not
            ignoring any bboxes.
        ignore_wrt_candidates (bool): Whether to compute the iof between
            `bboxes` and `gt_bboxes_ignore`, or the contrary.
        match_low_quality (bool): Whether to allow low quality matches. This is
            usually allowed for RPN and single stage detectors, but not allowed
            in the second stage. Details are demonstrated in Step 4.
        gpu_assign_thr (int): The upper bound of the number of GT for GPU
            assign. When the number of gt is above this threshold, will assign
            on CPU device. Negative values mean not assign on CPU.
        iou_calculator (dict): Config of overlaps Calculator.
    """

    def __init__(self,
                 pos_iou_thr: float,
                 neg_iou_thr: Union[float, tuple],
                 min_pos_iou: float = .0,
                 gt_max_assign_all: bool = True,
                 ignore_iof_thr: float = -1,
                 ignore_wrt_candidates: bool = True,
                 match_low_quality: bool = True,
                 iou_calculator: dict = dict(type='BboxOverlaps2D')):
        self.pos_iou_thr = pos_iou_thr
        self.neg_iou_thr = neg_iou_thr
        self.min_pos_iou = min_pos_iou
        self.gt_max_assign_all = gt_max_assign_all
        self.ignore_iof_thr = ignore_iof_thr
        self.ignore_wrt_candidates = ignore_wrt_candidates
        self.match_low_quality = match_low_quality
        self.iou_calculator = TASK_UTILS.build(iou_calculator)

    def assign(self,
               pred_instances: InstanceData,
               gt_instances: InstanceData,
               gt_instances_ignore: Optional[InstanceData] = None,
               **kwargs) -> AssignResult:
        """Assign gt to bboxes.

        This method assign a gt bbox to every bbox (proposal/anchor), each bbox
        will be assigned with -1, or a semi-positive number. -1 means negative
        sample, semi-positive number is the index (0-based) of assigned gt.
        The assignment is done in following steps, the order matters.

        1. assign every bbox to the background
        2. assign proposals whose iou with all gts < neg_iou_thr to 0
        3. for each bbox, if the iou with its nearest gt >= pos_iou_thr,
           assign it to that bbox
        4. for each gt bbox, assign its nearest proposals (may be more than
           one) to itself

        Args:
            pred_instances (:obj:`InstanceData`): Instances of model
                predictions. It includes ``priors``, and the priors can
                be anchors or points, or the bboxes predicted by the
                previous stage, has shape (n, 4). The bboxes predicted by
                the current model or stage will be named ``bboxes``,
                ``labels``, and ``scores``, the same as the ``InstanceData``
                in other places.
            gt_instances (:obj:`InstanceData`): Ground truth of instance
                annotations. It usually includes ``bboxes``, with shape (k, 4),
                and ``labels``, with shape (k, ).
            gt_instances_ignore (:obj:`InstanceData`, optional): Instances
                to be ignored during training. It includes ``bboxes``
                attribute data that is ignored during training and testing.
                Defaults to None.

        Returns:
            :obj:`AssignResult`: The assign result.

        Example:
            >>> from mmengine.structures import InstanceData
            >>> self = MaxIoUAssigner(0.5, 0.5)
            >>> pred_instances = InstanceData()
            >>> pred_instances.priors = torch.Tensor([[0, 0, 10, 10],
            ...                                      [10, 10, 20, 20]])
            >>> gt_instances = InstanceData()
            >>> gt_instances.bboxes = torch.Tensor([[0, 0, 10, 9]])
            >>> gt_instances.labels = torch.Tensor([0])
            >>> assign_result = self.assign(pred_instances, gt_instances)
            >>> expected_gt_inds = torch.LongTensor([1, 0])
            >>> assert torch.all(assign_result.gt_inds == expected_gt_inds)
        """
        gt_bboxes = gt_instances.bboxes
        priors = pred_instances.priors
        gt_labels = gt_instances.labels
        if gt_instances_ignore is not None:
            gt_bboxes_ignore = gt_instances_ignore.bboxes
        else:
            gt_bboxes_ignore = None

        overlaps = self.iou_calculator(gt_bboxes, priors)

        if (self.ignore_iof_thr > 0 and gt_bboxes_ignore is not None
                and gt_bboxes_ignore.numel() > 0 and priors.numel() > 0):
            if self.ignore_wrt_candidates:
                ignore_overlaps = self.iou_calculator(
                    priors, gt_bboxes_ignore, mode='iof')
                ignore_max_overlaps, _ = ignore_overlaps.max(dim=1)
            else:
                ignore_overlaps = self.iou_calculator(
                    gt_bboxes_ignore, priors, mode='iof')
                ignore_max_overlaps, _ = ignore_overlaps.max(dim=0)
            overlaps[:, ignore_max_overlaps > self.ignore_iof_thr] = -1

        assign_result = self.assign_wrt_overlaps(overlaps, gt_labels)
        return assign_result

    def assign_wrt_overlaps(self, overlaps: jt.Var,
                            gt_labels: jt.Var) -> AssignResult:
        """Assign w.r.t. the overlaps of priors with gts.

        Args:
            overlaps (Tensor): Overlaps between k gt_bboxes and n bboxes,
                shape(k, n).
            gt_labels (Tensor): Labels of k gt_bboxes, shape (k, ).

        Returns:
            :obj:`AssignResult`: The assign result.
        """
        num_gts, num_bboxes = overlaps.size(0), overlaps.size(1)

        # 1. assign -1 by default
        assigned_gt_inds = jt.full((num_bboxes, ), -1, dtype=jt.int64)

        if num_gts == 0 or num_bboxes == 0:
            # No ground truth or boxes, return empty assignment
            max_overlaps = jt.zeros((num_bboxes, ), dtype=overlaps.dtype)
            assigned_labels = jt.full((num_bboxes, ), -1, dtype=jt.int64)
            if num_gts == 0:
                # No truth, assign everything to background
                assigned_gt_inds[:] = 0
            return AssignResult(
                num_gts=num_gts,
                gt_inds=assigned_gt_inds,
                max_overlaps=max_overlaps,
                labels=assigned_labels)

        # for each anchor, which gt best overlaps with it
        # for each anchor, the max iou of all gts
        argmax_overlaps, max_overlaps = jt.argmax(overlaps, dim=0)
        # for each gt, which anchor best overlaps with it
        # for each gt, the max iou of all proposals
        gt_argmax_overlaps, gt_max_overlaps = jt.argmax(overlaps, dim=1)

        # 2. assign negative: below
        # the negative inds are set to be 0
        if isinstance(self.neg_iou_thr, float):
            assigned_gt_inds[(max_overlaps >= 0)
                             & (max_overlaps < self.neg_iou_thr)] = 0
        elif isinstance(self.neg_iou_thr, tuple):
            assert len(self.neg_iou_thr) == 2
            assigned_gt_inds[(max_overlaps >= self.neg_iou_thr[0])
                             & (max_overlaps < self.neg_iou_thr[1])] = 0

        # 3. assign positive: above positive IoU threshold
        pos_inds = max_overlaps >= self.pos_iou_thr
        assigned_gt_inds[pos_inds] = argmax_overlaps[pos_inds] + 1

        if self.match_low_quality:
            # Low-quality matching will overwrite the assigned_gt_inds assigned
            # in Step 3. Thus, the assigned gt might not be the best one for
            # prediction.
            # For example, if bbox A has 0.9 and 0.8 iou with GT bbox 1 & 2,
            # bbox 1 will be assigned as the best target for bbox A in step 3.
            # However, if GT bbox 2's gt_argmax_overlaps = A, bbox A's
            # assigned_gt_inds will be overwritten to be bbox 2.
            # This might be the reason that it is not used in ROI Heads.
            for i in range(num_gts):
                if gt_max_overlaps[i] >= self.min_pos_iou:
                    if self.gt_max_assign_all:
                        max_iou_inds = overlaps[i, :] == gt_max_overlaps[i]
                        assigned_gt_inds[max_iou_inds] = i + 1
                    else:
                        assigned_gt_inds[gt_argmax_overlaps[i]] = i + 1

        assigned_labels = jt.full((num_bboxes, ),
                                  -1,
                                  dtype=assigned_gt_inds.dtype)
        pos_inds = jt.nonzero(assigned_gt_inds > 0).squeeze(dim=1)
        if pos_inds.numel() > 0:
            assigned_labels[pos_inds] = gt_labels[assigned_gt_inds[pos_inds] -
                                                  1]

        return AssignResult(
            num_gts=num_gts,
            gt_inds=assigned_gt_inds,
            max_overlaps=max_overlaps,
            labels=assigned_labels)

jittordet/models/task_utils/assigners/max_iou_assigner.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/coders/base_bbox_coder.py
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod


class BaseBBoxCoder(metaclass=ABCMeta):
    """Base bounding box coder.

    Args:
        use_box_type (bool): Whether to warp decoded boxes with the
            box type data structure. Defaults to False.
    """

    @abstractmethod
    def encode(self, bboxes, gt_bboxes):
        """Encode deltas between bboxes and ground truth boxes."""

    @abstractmethod
    def decode(self, bboxes, bboxes_pred):
        """Decode the predicted bboxes according to prediction and base
        boxes."""

jittordet/models/task_utils/bbox_coders/base_bbox_coder.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/coders/delta_xywh_bbox_coder.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt
import numpy as np

from jittordet.engine import TASK_UTILS
from .base_bbox_coder import BaseBBoxCoder


@TASK_UTILS.register_module()
class DeltaXYWHBBoxCoder(BaseBBoxCoder):
    """Delta XYWH BBox coder.

    Following the practice in `R-CNN <https://arxiv.org/abs/1311.2524>`_,
    this coder encodes bbox (x1, y1, x2, y2) into delta (dx, dy, dw, dh) and
    decodes delta (dx, dy, dw, dh) back to original bbox (x1, y1, x2, y2).

    Args:
        target_means (Sequence[float]): Denormalizing means of target for
            delta coordinates
        target_stds (Sequence[float]): Denormalizing standard deviation of
            target for delta coordinates
        clip_border (bool, optional): Whether clip the objects outside the
            border of the image. Defaults to True.
        add_ctr_clamp (bool): Whether to add center clamp, when added, the
            predicted box is clamped is its center is too far away from
            the original anchor's center. Only used by YOLOF. Default False.
        ctr_clamp (int): the maximum pixel shift to clamp. Only used by YOLOF.
            Default 32.
    """

    def __init__(self,
                 target_means=(0., 0., 0., 0.),
                 target_stds=(1., 1., 1., 1.),
                 clip_border=True,
                 add_ctr_clamp=False,
                 ctr_clamp=32,
                 **kwargs):
        super().__init__(**kwargs)
        self.means = target_means
        self.stds = target_stds
        self.clip_border = clip_border
        self.add_ctr_clamp = add_ctr_clamp
        self.ctr_clamp = ctr_clamp

    def encode(self, bboxes, gt_bboxes):
        """Get box regression transformation deltas that can be used to
        transform the ``bboxes`` into the ``gt_bboxes``.

        Args:
            bboxes (torch.Tensor or :obj:`BaseBoxes`): Source boxes,
                e.g., object proposals.
            gt_bboxes (torch.Tensor or :obj:`BaseBoxes`): Target of the
                transformation, e.g., ground-truth boxes.

        Returns:
            torch.Tensor: Box transformation deltas
        """
        assert bboxes.size(0) == gt_bboxes.size(0)
        assert bboxes.size(-1) == gt_bboxes.size(-1) == 4
        encoded_bboxes = bbox2delta(bboxes, gt_bboxes, self.means, self.stds)
        return encoded_bboxes

    def decode(self,
               bboxes,
               pred_bboxes,
               max_shape=None,
               wh_ratio_clip=16 / 1000):
        """Apply transformation `pred_bboxes` to `boxes`.

        Args:
            bboxes (torch.Tensor or :obj:`BaseBoxes`): Basic boxes. Shape
                (B, N, 4) or (N, 4)
            pred_bboxes (Tensor): Encoded offsets with respect to each roi.
               Has shape (B, N, num_classes * 4) or (B, N, 4) or
               (N, num_classes * 4) or (N, 4). Note N = num_anchors * W * H
               when rois is a grid of anchors.Offset encoding follows [1]_.
            max_shape (Sequence[int] or torch.Tensor or Sequence[
               Sequence[int]],optional): Maximum bounds for boxes, specifies
               (H, W, C) or (H, W). If bboxes shape is (B, N, 4), then
               the max_shape should be a Sequence[Sequence[int]]
               and the length of max_shape should also be B.
            wh_ratio_clip (float, optional): The allowed ratio between
                width and height.

        Returns:
            Union[torch.Tensor, :obj:`BaseBoxes`]: Decoded boxes.
        """
        assert pred_bboxes.size(0) == bboxes.size(0)
        if pred_bboxes.ndim == 3:
            assert pred_bboxes.size(1) == bboxes.size(1)

        decoded_bboxes = delta2bbox(bboxes, pred_bboxes, self.means, self.stds,
                                    max_shape, wh_ratio_clip, self.clip_border,
                                    self.add_ctr_clamp, self.ctr_clamp)

        return decoded_bboxes


def bbox2delta(proposals, gt, means=(0., 0., 0., 0.), stds=(1., 1., 1., 1.)):
    """Compute deltas of proposals w.r.t. gt.

    We usually compute the deltas of x, y, w, h of proposals w.r.t ground
    truth bboxes to get regression target.
    This is the inverse function of :func:`delta2bbox`.

    Args:
        proposals (Tensor): Boxes to be transformed, shape (N, ..., 4)
        gt (Tensor): Gt bboxes to be used as base, shape (N, ..., 4)
        means (Sequence[float]): Denormalizing means for delta coordinates
        stds (Sequence[float]): Denormalizing standard deviation for delta
            coordinates

    Returns:
        Tensor: deltas with shape (N, 4), where columns represent dx, dy,
            dw, dh.
    """
    assert proposals.size() == gt.size()

    proposals = proposals.float()
    gt = gt.float()
    px = (proposals[..., 0] + proposals[..., 2]) * 0.5
    py = (proposals[..., 1] + proposals[..., 3]) * 0.5
    pw = proposals[..., 2] - proposals[..., 0]
    ph = proposals[..., 3] - proposals[..., 1]

    gx = (gt[..., 0] + gt[..., 2]) * 0.5
    gy = (gt[..., 1] + gt[..., 3]) * 0.5
    gw = gt[..., 2] - gt[..., 0]
    gh = gt[..., 3] - gt[..., 1]

    dx = (gx - px) / pw
    dy = (gy - py) / ph
    dw = jt.log(gw / pw)
    dh = jt.log(gh / ph)
    deltas = jt.stack([dx, dy, dw, dh], dim=-1)

    means = jt.array(means, dtype=deltas.dtype).unsqueeze(0)
    stds = jt.array(stds, dtype=deltas.dtype).unsqueeze(0)
    deltas = (deltas - means) / stds
    return deltas


def delta2bbox(rois,
               deltas,
               means=(0., 0., 0., 0.),
               stds=(1., 1., 1., 1.),
               max_shape=None,
               wh_ratio_clip=16 / 1000,
               clip_border=True,
               add_ctr_clamp=False,
               ctr_clamp=32):
    """Apply deltas to shift/scale base boxes.

    Typically the rois are anchor or proposed bounding boxes and the deltas are
    network outputs used to shift/scale those boxes.
    This is the inverse function of :func:`bbox2delta`.

    Args:
        rois (Tensor): Boxes to be transformed. Has shape (N, 4).
        deltas (Tensor): Encoded offsets relative to each roi.
            Has shape (N, num_classes * 4) or (N, 4). Note
            N = num_base_anchors * W * H, when rois is a grid of
            anchors. Offset encoding follows [1]_.
        means (Sequence[float]): Denormalizing means for delta coordinates.
            Default (0., 0., 0., 0.).
        stds (Sequence[float]): Denormalizing standard deviation for delta
            coordinates. Default (1., 1., 1., 1.).
        max_shape (tuple[int, int]): Maximum bounds for boxes, specifies
           (H, W). Default None.
        wh_ratio_clip (float): Maximum aspect ratio for boxes. Default
            16 / 1000.
        clip_border (bool, optional): Whether clip the objects outside the
            border of the image. Default True.
        add_ctr_clamp (bool): Whether to add center clamp. When set to True,
            the center of the prediction bounding box will be clamped to
            avoid being too far away from the center of the anchor.
            Only used by YOLOF. Default False.
        ctr_clamp (int): the maximum pixel shift to clamp. Only used by YOLOF.
            Default 32.

    Returns:
        Tensor: Boxes with shape (N, num_classes * 4) or (N, 4), where 4
           represent tl_x, tl_y, br_x, br_y.

    References:
        .. [1] https://arxiv.org/abs/1311.2524

    Example:
        >>> rois = torch.Tensor([[ 0.,  0.,  1.,  1.],
        >>>                      [ 0.,  0.,  1.,  1.],
        >>>                      [ 0.,  0.,  1.,  1.],
        >>>                      [ 5.,  5.,  5.,  5.]])
        >>> deltas = torch.Tensor([[  0.,   0.,   0.,   0.],
        >>>                        [  1.,   1.,   1.,   1.],
        >>>                        [  0.,   0.,   2.,  -1.],
        >>>                        [ 0.7, -1.9, -0.5,  0.3]])
        >>> delta2bbox(rois, deltas, max_shape=(32, 32, 3))
        tensor([[0.0000, 0.0000, 1.0000, 1.0000],
                [0.1409, 0.1409, 2.8591, 2.8591],
                [0.0000, 0.3161, 4.1945, 0.6839],
                [5.0000, 5.0000, 5.0000, 5.0000]])
    """
    num_bboxes, num_classes = deltas.size(0), deltas.size(1) // 4
    if num_bboxes == 0:
        return deltas

    deltas = deltas.reshape(-1, 4)

    means = jt.array(means, dtype=deltas.dtype).view(1, -1)
    stds = jt.array(stds, dtype=deltas.dtype).view(1, -1)
    denorm_deltas = deltas * stds + means

    dxy = denorm_deltas[:, :2]
    dwh = denorm_deltas[:, 2:]

    # Compute width/height of each roi
    rois_ = rois.repeat(1, num_classes).reshape(-1, 4)
    pxy = ((rois_[:, :2] + rois_[:, 2:]) * 0.5)
    pwh = (rois_[:, 2:] - rois_[:, :2])

    dxy_wh = pwh * dxy

    max_ratio = np.abs(np.log(wh_ratio_clip))
    if add_ctr_clamp:
        dxy_wh = jt.clamp(dxy_wh, max_v=ctr_clamp, min_v=-ctr_clamp)
        dwh = jt.clamp(dwh, max_v=max_ratio)
    else:
        dwh = dwh.clamp(min_v=-max_ratio, max_v=max_ratio)

    gxy = pxy + dxy_wh
    gwh = pwh * dwh.exp()
    x1y1 = gxy - (gwh * 0.5)
    x2y2 = gxy + (gwh * 0.5)
    bboxes = jt.concat([x1y1, x2y2], dim=-1)
    if clip_border and max_shape is not None:
        bboxes[..., 0::2].clamp_(min_v=0, max_v=max_shape[1])
        bboxes[..., 1::2].clamp_(min_v=0, max_v=max_shape[0])
    bboxes = bboxes.reshape(num_bboxes, -1)
    return bboxes

jittordet/models/task_utils/bbox_coders/delta_xywh_bbox_coder.py
JittorCV/jittordet
from .anchor_generator import AnchorGenerator
from .utils import anchor_inside_flags

__all__ = ['AnchorGenerator', 'anchor_inside_flags']

jittordet/models/task_utils/prior_generators/__init__.py
JittorCV/jittordet
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt
import numpy as np

from jittordet.engine import TASK_UTILS


@TASK_UTILS.register_module()
class AnchorGenerator:
    """Standard anchor generator for 2D anchor-based detectors.

    Args:
        strides (list[int] | list[tuple[int, int]]): Strides of anchors
            in multiple feature levels in order (w, h).
        ratios (list[float]): The list of ratios between the height and width
            of anchors in a single level.
        scales (list[int] | None): Anchor scales for anchors in a single level.
            It cannot be set at the same time if `octave_base_scale` and
            `scales_per_octave` are set.
        base_sizes (list[int] | None): The basic sizes
            of anchors in multiple levels.
            If None is given, strides will be used as base_sizes.
            (If strides are non square, the shortest stride is taken.)
        scale_major (bool): Whether to multiply scales first when generating
            base anchors. If true, the anchors in the same row will have the
            same scales. By default it is True in V2.0
        octave_base_scale (int): The base scale of octave.
        scales_per_octave (int): Number of scales for each octave.
            `octave_base_scale` and `scales_per_octave` are usually used in
            retinanet and the `scales` should be None when they are set.
        centers (list[tuple[float, float]] | None): The centers of the anchor
            relative to the feature grid center in multiple feature levels.
            By default it is set to be None and not used. If a list of tuple of
            float is given, they will be used to shift the centers of anchors.
        center_offset (float): The offset of center in proportion to anchors'
            width and height. By default it is 0 in V2.0.

    Examples:
        >>> from mmdet.core import AnchorGenerator
        >>> self = AnchorGenerator([16], [1.], [1.], [9])
        >>> all_anchors = self.grid_priors([(2, 2)], device='cpu')
        >>> print(all_anchors)
        [tensor([[-4.5000, -4.5000,  4.5000,  4.5000],
                [11.5000, -4.5000, 20.5000,  4.5000],
                [-4.5000, 11.5000,  4.5000, 20.5000],
                [11.5000, 11.5000, 20.5000, 20.5000]])]
        >>> self = AnchorGenerator([16, 32], [1.], [1.], [9, 18])
        >>> all_anchors = self.grid_priors([(2, 2), (1, 1)], device='cpu')
        >>> print(all_anchors)
        [tensor([[-4.5000, -4.5000,  4.5000,  4.5000],
                [11.5000, -4.5000, 20.5000,  4.5000],
                [-4.5000, 11.5000,  4.5000, 20.5000],
                [11.5000, 11.5000, 20.5000, 20.5000]]), \
        tensor([[-9., -9., 9., 9.]])]
    """

    def __init__(self,
                 strides,
                 ratios,
                 scales=None,
                 base_sizes=None,
                 scale_major=True,
                 octave_base_scale=None,
                 scales_per_octave=None,
                 centers=None,
                 center_offset=0.):
        # check center and center_offset
        if center_offset != 0:
            assert centers is None, 'center cannot be set when center_offset' \
                                    f'!=0, {centers} is given.'
        if not (0 <= center_offset <= 1):
            raise ValueError('center_offset should be in range [0, 1], '
                             f'{center_offset} is given.')
        if centers is not None:
            assert len(centers) == len(strides), \
                'The number of strides should be the same as centers, got ' \
                f'{strides} and {centers}'

        # calculate base sizes of anchors
        self.strides = [tuple([stride, stride]) for stride in strides]
        self.base_sizes = [min(stride) for stride in self.strides
                           ] if base_sizes is None else base_sizes
        assert len(self.base_sizes) == len(self.strides), \
            'The number of strides should be the same as base sizes, got ' \
            f'{self.strides} and {self.base_sizes}'

        # calculate scales of anchors
        assert ((octave_base_scale is not None
                 and scales_per_octave is not None) ^ (scales is not None)), \
            'scales and octave_base_scale with scales_per_octave cannot' \
            ' be set at the same time'
        if scales is not None:
            self.scales = jt.array(scales)
        elif octave_base_scale is not None and scales_per_octave is not None:
            octave_scales = np.array(
                [2**(i / scales_per_octave) for i in range(scales_per_octave)])
            scales = octave_scales * octave_base_scale
            self.scales = jt.array(scales)
        else:
            raise ValueError('Either scales or octave_base_scale with '
                             'scales_per_octave should be set')

        self.octave_base_scale = octave_base_scale
        self.scales_per_octave = scales_per_octave
        self.ratios = jt.array(ratios)
        self.scale_major = scale_major
        self.centers = centers
        self.center_offset = center_offset
        self.base_anchors = self.gen_base_anchors()

    @property
    def num_base_priors(self):
        """list[int]: The number of priors (anchors) at a point
        on the feature grid"""
        return [base_anchors.size(0) for base_anchors in self.base_anchors]

    @property
    def num_levels(self):
        """int: number of feature levels that the generator will be applied"""
        return len(self.strides)

    def gen_base_anchors(self):
        """Generate base anchors.

        Returns:
            list(torch.Tensor): Base anchors of a feature grid in multiple \
                feature levels.
        """
        multi_level_base_anchors = []
        for i, base_size in enumerate(self.base_sizes):
            center = None
            if self.centers is not None:
                center = self.centers[i]
            multi_level_base_anchors.append(
                self.gen_single_level_base_anchors(
                    base_size,
                    scales=self.scales,
                    ratios=self.ratios,
                    center=center))
        return multi_level_base_anchors

    def gen_single_level_base_anchors(self,
                                      base_size,
                                      scales,
                                      ratios,
                                      center=None):
        """Generate base anchors of a single level.

        Args:
            base_size (int | float): Basic size of an anchor.
            scales (torch.Tensor): Scales of the anchor.
            ratios (torch.Tensor): The ratio between between the height
                and width of anchors in a single level.
            center (tuple[float], optional): The center of the base anchor
                related to a single feature grid. Defaults to None.

        Returns:
            torch.Tensor: Anchors in a single-level feature maps.
        """
        w = base_size
        h = base_size
        if center is None:
            x_center = self.center_offset * w
            y_center = self.center_offset * h
        else:
            x_center, y_center = center

        h_ratios = jt.sqrt(ratios)
        w_ratios = 1 / h_ratios
        if self.scale_major:
            ws = (w * w_ratios[:, None] * scales[None, :]).view(-1)
            hs = (h * h_ratios[:, None] * scales[None, :]).view(-1)
        else:
            ws = (w * scales[:, None] * w_ratios[None, :]).view(-1)
            hs = (h * scales[:, None] * h_ratios[None, :]).view(-1)

        # use float anchor and the anchor's center is aligned with the
        # pixel center
        base_anchors = [
            x_center - 0.5 * ws, y_center - 0.5 * hs, x_center + 0.5 * ws,
            y_center + 0.5 * hs
        ]
        base_anchors = jt.stack(base_anchors, dim=-1)

        return base_anchors

    def _meshgrid(self, x, y, row_major=True):
        """Generate mesh grid of x and y.

        Args:
            x (torch.Tensor): Grids of x dimension.
            y (torch.Tensor): Grids of y dimension.
            row_major (bool, optional): Whether to return y grids first.
                Defaults to True.

        Returns:
            tuple[torch.Tensor]: The mesh grids of x and y.
        """
        # use shape instead of len to keep tracing while exporting to onnx
        xx = x.repeat(y.shape[0])
        yy = y.view(-1, 1).repeat(1, x.shape[0]).view(-1)
        if row_major:
            return xx, yy
        else:
            return yy, xx

    def grid_priors(self, featmap_sizes, dtype=jt.float32):
        """Generate grid anchors in multiple feature levels.

        Args:
            featmap_sizes (list[tuple]): List of feature map sizes in
                multiple feature levels.
            dtype (:obj:`torch.dtype`): Dtype of priors.
                Default: torch.float32.
            device (str): The device where the anchors will be put on.

        Return:
            list[torch.Tensor]: Anchors in multiple feature levels. \
                The sizes of each tensor should be [N, 4], where \
                N = width * height * num_base_anchors, width and height \
                are the sizes of the corresponding feature level, \
                num_base_anchors is the number of anchors for that level.
        """
        assert self.num_levels == len(featmap_sizes)
        multi_level_anchors = []
        for i in range(self.num_levels):
            anchors = self.single_level_grid_priors(
                featmap_sizes[i], level_idx=i, dtype=dtype)
            multi_level_anchors.append(anchors)
        return multi_level_anchors

    def single_level_grid_priors(self,
                                 featmap_size,
                                 level_idx,
                                 dtype=jt.float32):
        """Generate grid anchors of a single level.

        Note:
            This function is usually called by method ``self.grid_priors``.

        Args:
            featmap_size (tuple[int]): Size of the feature maps.
            level_idx (int): The index of corresponding feature map level.
            dtype (obj:`torch.dtype`): Date type of points.Defaults to
                ``torch.float32``.
            device (str, optional): The device the tensor will be put on.
                Defaults to 'cuda'.

        Returns:
            torch.Tensor: Anchors in the overall feature maps.
        """

        base_anchors = self.base_anchors[level_idx].astype(dtype)
        feat_h, feat_w = featmap_size
        stride_w, stride_h = self.strides[level_idx]
        # First create Range with the default dtype, than convert to
        # target `dtype` for onnx exporting.

        shift_x = jt.arange(0, feat_w).astype(dtype) * stride_w
        shift_y = jt.arange(0, feat_h).astype(dtype) * stride_h

        shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
        shifts = jt.stack([shift_xx, shift_yy, shift_xx, shift_yy], dim=-1)
        # first feat_w elements correspond to the first row of shifts
        # add A anchors (1, A, 4) to K shifts (K, 1, 4) to get
        # shifted anchors (K, A, 4), reshape to (K*A, 4)

        all_anchors = base_anchors[None, :, :] + shifts[:, None, :]
        all_anchors = all_anchors.view(-1, 4)
        # first A rows correspond to A anchors of (0, 0) in feature map,
        # then (0, 1), (0, 2), ...
        return all_anchors

    def sparse_priors(self,
                      prior_idxs,
                      featmap_size,
                      level_idx,
                      dtype=jt.float32):
        """Generate sparse anchors according to the ``prior_idxs``.

        Args:
            prior_idxs (Tensor): The index of corresponding anchors
                in the feature map.
            featmap_size (tuple[int]): feature map size arrange as (h, w).
            level_idx (int): The level index of corresponding feature
                map.
            dtype (obj:`torch.dtype`): Date type of points.Defaults to
                ``torch.float32``.
            device (obj:`torch.device`): The device where the points is
                located.
        Returns:
            Tensor: Anchor with shape (N, 4), N should be equal to
                the length of ``prior_idxs``.
        """

        height, width = featmap_size
        num_base_anchors = self.num_base_priors[level_idx]
        base_anchor_id = prior_idxs % num_base_anchors
        x = (prior_idxs //
             num_base_anchors) % width * self.strides[level_idx][0]
        y = (prior_idxs // width //
             num_base_anchors) % height * self.strides[level_idx][1]
        priors = jt.stack([x, y, x, y], 1).astype(dtype) + \
            self.base_anchors[level_idx][base_anchor_id, :]

        return priors

    def valid_flags(self, featmap_sizes, pad_shape):
        """Generate valid flags of anchors in multiple feature levels.

        Args:
            featmap_sizes (list(tuple)): List of feature map sizes in
                multiple feature levels.
            pad_shape (tuple): The padded shape of the image.
            device (str): Device where the anchors will be put on.

        Return:
            list(torch.Tensor): Valid flags of anchors in multiple levels.
        """
        assert self.num_levels == len(featmap_sizes)
        multi_level_flags = []
        for i in range(self.num_levels):
            anchor_stride = self.strides[i]
            feat_h, feat_w = featmap_sizes[i]
            h, w = pad_shape[:2]
            valid_feat_h = min(int(np.ceil(h / anchor_stride[1])), feat_h)
            valid_feat_w = min(int(np.ceil(w / anchor_stride[0])), feat_w)
            flags = self.single_level_valid_flags((feat_h, feat_w),
                                                  (valid_feat_h, valid_feat_w),
                                                  self.num_base_priors[i])
            multi_level_flags.append(flags)
        return multi_level_flags

    def single_level_valid_flags(self, featmap_size, valid_size,
                                 num_base_anchors):
        """Generate the valid flags of anchor in a single feature map.

        Args:
            featmap_size (tuple[int]): The size of feature maps, arrange
                as (h, w).
            valid_size (tuple[int]): The valid size of the feature maps.
            num_base_anchors (int): The number of base anchors.
            device (str, optional): Device where the flags will be put on.
                Defaults to 'cuda'.

        Returns:
            torch.Tensor: The valid flags of each anchor in a single level \
                feature map.
        """
        feat_h, feat_w = featmap_size
        valid_h, valid_w = valid_size
        assert valid_h <= feat_h and valid_w <= feat_w
        valid_x = jt.zeros(feat_w, dtype=jt.bool)
        valid_y = jt.zeros(feat_h, dtype=jt.bool)
        valid_x[:valid_w] = 1
        valid_y[:valid_h] = 1
        valid_xx, valid_yy = self._meshgrid(valid_x, valid_y)
        valid = valid_xx & valid_yy
        valid = valid[:, None].expand(valid.size(0),
                                      num_base_anchors).contiguous().view(-1)
        return valid

    def __repr__(self):
        """str: a string that describes the module"""
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
        repr_str += f'{indent_str}strides={self.strides},\n'
        repr_str += f'{indent_str}ratios={self.ratios},\n'
        repr_str += f'{indent_str}scales={self.scales},\n'
        repr_str += f'{indent_str}base_sizes={self.base_sizes},\n'
        repr_str += f'{indent_str}scale_major={self.scale_major},\n'
        repr_str += f'{indent_str}octave_base_scale='
        repr_str += f'{self.octave_base_scale},\n'
        repr_str += f'{indent_str}scales_per_octave='
        repr_str += f'{self.scales_per_octave},\n'
        repr_str += f'{indent_str}num_levels={self.num_levels}\n'
        repr_str += f'{indent_str}centers={self.centers},\n'
        repr_str += f'{indent_str}center_offset={self.center_offset})'
        return repr_str

jittordet/models/task_utils/prior_generators/anchor_generator.py
JittorCV/jittordet
def anchor_inside_flags(flat_anchors,
                        valid_flags,
                        img_shape,
                        allowed_border=0):
    """Check whether the anchors are inside the border.

    Args:
        flat_anchors (torch.Tensor): Flatten anchors, shape (n, 4).
        valid_flags (torch.Tensor): An existing valid flags of anchors.
        img_shape (tuple(int)): Shape of current image.
        allowed_border (int, optional): The border to allow the valid anchor.
            Defaults to 0.

    Returns:
        torch.Tensor: Flags indicating whether the anchors are inside a \
            valid range.
    """
    img_h, img_w = img_shape[:2]
    if allowed_border >= 0:
        inside_flags = valid_flags & \
            (flat_anchors[:, 0] >= -allowed_border) & \
            (flat_anchors[:, 1] >= -allowed_border) & \
            (flat_anchors[:, 2] < img_w + allowed_border) & \
            (flat_anchors[:, 3] < img_h + allowed_border)
    else:
        inside_flags = valid_flags
    return inside_flags

jittordet/models/task_utils/prior_generators/utils.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/samplers/base_sampler.py
# Copyright (c) OpenMMLab. All rights reserved.
from abc import ABCMeta, abstractmethod

import jittor as jt

from jittordet.structures import InstanceData
from ..assigners import AssignResult
from .sampling_result import SamplingResult


class BaseSampler(metaclass=ABCMeta):
    """Base class of samplers.

    Args:
        num (int): Number of samples
        pos_fraction (float): Fraction of positive samples
        neg_pos_up (int): Upper bound number of negative and
            positive samples. Defaults to -1.
        add_gt_as_proposals (bool): Whether to add ground truth
            boxes as proposals. Defaults to True.
    """

    def __init__(self,
                 num: int,
                 pos_fraction: float,
                 neg_pos_ub: int = -1,
                 add_gt_as_proposals: bool = True,
                 **kwargs) -> None:
        self.num = num
        self.pos_fraction = pos_fraction
        self.neg_pos_ub = neg_pos_ub
        self.add_gt_as_proposals = add_gt_as_proposals
        self.pos_sampler = self
        self.neg_sampler = self

    @abstractmethod
    def _sample_pos(self, assign_result: AssignResult, num_expected: int,
                    **kwargs):
        """Sample positive samples."""
        pass

    @abstractmethod
    def _sample_neg(self, assign_result: AssignResult, num_expected: int,
                    **kwargs):
        """Sample negative samples."""
        pass

    def sample(self, assign_result: AssignResult, pred_instances: InstanceData,
               gt_instances: InstanceData, **kwargs) -> SamplingResult:
        """Sample positive and negative bboxes.

        This is a simple implementation of bbox sampling given candidates,
        assigning results and ground truth bboxes.

        Args:
            assign_result (:obj:`AssignResult`): Assigning results.
            pred_instances (:obj:`InstanceData`): Instances of model
                predictions. It includes ``priors``, and the priors can
                be anchors or points, or the bboxes predicted by the
                previous stage, has shape (n, 4). The bboxes predicted by
                the current model or stage will be named ``bboxes``,
                ``labels``, and ``scores``, the same as the ``InstanceData``
                in other places.
            gt_instances (:obj:`InstanceData`): Ground truth of instance
                annotations. It usually includes ``bboxes``, with shape (k, 4),
                and ``labels``, with shape (k, ).

        Returns:
            :obj:`SamplingResult`: Sampling result.

        Example:
            >>> from mmengine.structures import InstanceData
            >>> from mmdet.models.task_modules.samplers import RandomSampler,
            >>> from mmdet.models.task_modules.assigners import AssignResult
            >>> from mmdet.models.task_modules.samplers.
            ... sampling_result import ensure_rng, random_boxes
            >>> rng = ensure_rng(None)
            >>> assign_result = AssignResult.random(rng=rng)
            >>> pred_instances = InstanceData()
            >>> pred_instances.priors = random_boxes(assign_result.num_preds,
            ...                                      rng=rng)
            >>> gt_instances = InstanceData()
            >>> gt_instances.bboxes = random_boxes(assign_result.num_gts,
            ...                                    rng=rng)
            >>> gt_instances.labels = torch.randint(
            ...     0, 5, (assign_result.num_gts,), dtype=torch.long)
            >>> self = RandomSampler(num=32, pos_fraction=0.5, neg_pos_ub=-1,
            >>>                      add_gt_as_proposals=False)
            >>> self = self.sample(assign_result, pred_instances, gt_instances)
        """
        gt_bboxes = gt_instances.bboxes
        priors = pred_instances.priors
        gt_labels = gt_instances.labels
        if len(priors.shape) < 2:
            priors = priors[None, :]

        gt_flags = jt.zeros((priors.shape[0], ), dtype=jt.uint8)
        if self.add_gt_as_proposals and len(gt_bboxes) > 0:
            priors = jt.concat([gt_bboxes, priors], dim=0)
            assign_result.add_gt_(gt_labels)
            gt_ones = jt.ones(gt_bboxes.shape[0], dtype=jt.uint8)
            gt_flags = jt.concat([gt_ones, gt_flags])

        num_expected_pos = int(self.num * self.pos_fraction)
        pos_inds = self.pos_sampler._sample_pos(
            assign_result, num_expected_pos, bboxes=priors, **kwargs)
        # We found that sampled indices have duplicated items occasionally.
        # (may be a bug of PyTorch)
        pos_inds = pos_inds.unique()
        num_sampled_pos = pos_inds.numel()
        num_expected_neg = self.num - num_sampled_pos
        if self.neg_pos_ub >= 0:
            _pos = max(1, num_sampled_pos)
            neg_upper_bound = int(self.neg_pos_ub * _pos)
            if num_expected_neg > neg_upper_bound:
                num_expected_neg = neg_upper_bound
        neg_inds = self.neg_sampler._sample_neg(
            assign_result, num_expected_neg, bboxes=priors, **kwargs)
        neg_inds = neg_inds.unique()

        sampling_result = SamplingResult(
            pos_inds=pos_inds,
            neg_inds=neg_inds,
            priors=priors,
            gt_bboxes=gt_bboxes,
            assign_result=assign_result,
            gt_flags=gt_flags)
        return sampling_result

jittordet/models/task_utils/samplers/base_sampler.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/samplers/pseudo_sampler.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt

from jittordet.engine import TASK_UTILS
from jittordet.structures import InstanceData
from ..assigners import AssignResult
from .base_sampler import BaseSampler
from .sampling_result import SamplingResult


@TASK_UTILS.register_module()
class PseudoSampler(BaseSampler):
    """A pseudo sampler that does not do sampling actually."""

    def __init__(self, **kwargs):
        pass

    def _sample_pos(self, **kwargs):
        """Sample positive samples."""
        raise NotImplementedError

    def _sample_neg(self, **kwargs):
        """Sample negative samples."""
        raise NotImplementedError

    def sample(self, assign_result: AssignResult, pred_instances: InstanceData,
               gt_instances: InstanceData, *args, **kwargs):
        """Directly returns the positive and negative indices  of samples.

        Args:
            assign_result (:obj:`AssignResult`): Bbox assigning results.
            pred_instances (:obj:`InstanceData`): Instances of model
                predictions. It includes ``priors``, and the priors can
                be anchors, points, or bboxes predicted by the model,
                shape(n, 4).
            gt_instances (:obj:`InstanceData`): Ground truth of instance
                annotations. It usually includes ``bboxes`` and ``labels``
                attributes.

        Returns:
            :obj:`SamplingResult`: sampler results
        """
        gt_bboxes = gt_instances.bboxes
        priors = pred_instances.priors

        pos_inds = jt.nonzero(assign_result.gt_inds > 0).squeeze(-1).unique()
        neg_inds = jt.nonzero(assign_result.gt_inds == 0).squeeze(-1).unique()

        gt_flags = jt.zeros(priors.shape[0], dtype=jt.uint8)
        sampling_result = SamplingResult(
            pos_inds=pos_inds,
            neg_inds=neg_inds,
            priors=priors,
            gt_bboxes=gt_bboxes,
            assign_result=assign_result,
            gt_flags=gt_flags,
            avg_factor_with_neg=False)
        return sampling_result

jittordet/models/task_utils/samplers/pseudo_sampler.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/samplers/random_sampler.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Union

import jittor as jt
from numpy import ndarray

from jittordet.engine import TASK_UTILS
from ..assigners import AssignResult
from .base_sampler import BaseSampler


@TASK_UTILS.register_module()
class RandomSampler(BaseSampler):
    """Random sampler.

    Args:
        num (int): Number of samples
        pos_fraction (float): Fraction of positive samples
        neg_pos_up (int): Upper bound number of negative and
            positive samples. Defaults to -1.
        add_gt_as_proposals (bool): Whether to add ground truth
            boxes as proposals. Defaults to True.
    """

    def __init__(self,
                 num: int,
                 pos_fraction: float,
                 neg_pos_ub: int = -1,
                 add_gt_as_proposals: bool = True,
                 **kwargs):
        from jittordet.utils import ensure_rng
        super().__init__(
            num=num,
            pos_fraction=pos_fraction,
            neg_pos_ub=neg_pos_ub,
            add_gt_as_proposals=add_gt_as_proposals)
        self.rng = ensure_rng(kwargs.get('rng', None))

    def random_choice(self, gallery: Union[jt.Var, ndarray, list],
                      num: int) -> Union[jt.Var, ndarray]:
        """Random select some elements from the gallery.

        If `gallery` is a Tensor, the returned indices will be a Tensor;
        If `gallery` is a ndarray or list, the returned indices will be a
        ndarray.

        Args:
            gallery (Tensor | ndarray | list): indices pool.
            num (int): expected sample num.

        Returns:
            Tensor or ndarray: sampled indices.
        """
        assert len(gallery) >= num

        is_tensor = isinstance(gallery, jt.Var)
        if not is_tensor:
            gallery = jt.array(gallery, dtype=jt.int64)
        # This is a temporary fix. We can revert the following code
        # when PyTorch fixes the abnormal return of torch.randperm.
        # See: https://github.com/open-mmlab/mmdetection/pull/5014
        perm = jt.randperm(gallery.numel())[:num]
        rand_inds = gallery[perm]
        if not is_tensor:
            rand_inds = rand_inds.numpy()
        return rand_inds

    def _sample_pos(self, assign_result: AssignResult, num_expected: int,
                    **kwargs) -> Union[jt.Var, ndarray]:
        """Randomly sample some positive samples.

        Args:
            assign_result (:obj:`AssignResult`): Bbox assigning results.
            num_expected (int): The number of expected positive samples

        Returns:
            Tensor or ndarray: sampled indices.
        """
        pos_inds = jt.nonzero(assign_result.gt_inds > 0)
        if pos_inds.numel() != 0:
            pos_inds = pos_inds.squeeze(1)
        if pos_inds.numel() <= num_expected:
            return pos_inds
        else:
            return self.random_choice(pos_inds, num_expected)

    def _sample_neg(self, assign_result: AssignResult, num_expected: int,
                    **kwargs) -> Union[jt.Var, ndarray]:
        """Randomly sample some negative samples.

        Args:
            assign_result (:obj:`AssignResult`): Bbox assigning results.
            num_expected (int): The number of expected positive samples

        Returns:
            Tensor or ndarray: sampled indices.
        """
        neg_inds = jt.nonzero(assign_result.gt_inds == 0)
        if neg_inds.numel() != 0:
            neg_inds = neg_inds.squeeze(1)
        if len(neg_inds) <= num_expected:
            return neg_inds
        else:
            return self.random_choice(neg_inds, num_expected)

jittordet/models/task_utils/samplers/random_sampler.py
JittorCV/jittordet
# Modified from OpenMMLab.
# mmdet/models/task_modules/samplers/sampling_result.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt

from ..assigners import AssignResult


class SamplingResult:
    """Bbox sampling result.

    Args:
        pos_inds (Tensor): Indices of positive samples.
        neg_inds (Tensor): Indices of negative samples.
        priors (Tensor): The priors can be anchors or points,
            or the bboxes predicted by the previous stage.
        gt_bboxes (Tensor): Ground truth of bboxes.
        assign_result (:obj:`AssignResult`): Assigning results.
        gt_flags (Tensor): The Ground truth flags.
        avg_factor_with_neg (bool):  If True, ``avg_factor`` equal to
            the number of total priors; Otherwise, it is the number of
            positive priors. Defaults to True.

    Example:
        >>> # xdoctest: +IGNORE_WANT
        >>> from mmdet.models.task_modules.samplers.sampling_result import *  # NOQA
        >>> self = SamplingResult.random(rng=10)
        >>> print(f'self = {self}')
        self = <SamplingResult({
            'neg_inds': tensor([1,  2,  3,  5,  6,  7,  8,
                                9, 10, 11, 12, 13]),
            'neg_priors': torch.Size([12, 4]),
            'num_gts': 1,
            'num_neg': 12,
            'num_pos': 1,
            'avg_factor': 13,
            'pos_assigned_gt_inds': tensor([0]),
            'pos_inds': tensor([0]),
            'pos_is_gt': tensor([1], dtype=torch.uint8),
            'pos_priors': torch.Size([1, 4])
        })>
    """

    def __init__(self,
                 pos_inds: jt.Var,
                 neg_inds: jt.Var,
                 priors: jt.Var,
                 gt_bboxes: jt.Var,
                 assign_result: AssignResult,
                 gt_flags: jt.Var,
                 avg_factor_with_neg: bool = True) -> None:
        self.pos_inds = pos_inds
        self.neg_inds = neg_inds
        self.num_pos = max(pos_inds.numel(), 1)
        self.num_neg = max(neg_inds.numel(), 1)
        self.avg_factor_with_neg = avg_factor_with_neg
        self.avg_factor = self.num_pos + self.num_neg \
            if avg_factor_with_neg else self.num_pos
        self.pos_priors = priors[pos_inds]
        self.neg_priors = priors[neg_inds]
        self.pos_is_gt = gt_flags[pos_inds]

        self.num_gts = gt_bboxes.shape[0]
        self.pos_assigned_gt_inds = assign_result.gt_inds[pos_inds] - 1
        self.pos_gt_labels = assign_result.labels[pos_inds]
        if gt_bboxes.numel() == 0:
            # hack for index error case
            assert self.pos_assigned_gt_inds.numel() == 0
            self.pos_gt_bboxes = gt_bboxes.view(-1, 4)
        else:
            if len(gt_bboxes.shape) < 2:
                gt_bboxes = gt_bboxes.view(-1, 4)
            self.pos_gt_bboxes = gt_bboxes[self.pos_assigned_gt_inds.long()]

    @property
    def priors(self):
        """torch.Tensor: concatenated positive and negative priors"""
        return jt.concat([self.pos_priors, self.neg_priors])

    @property
    def info(self):
        """Returns a dictionary of info about the object."""
        return {
            'pos_inds': self.pos_inds,
            'neg_inds': self.neg_inds,
            'pos_priors': self.pos_priors,
            'neg_priors': self.neg_priors,
            'pos_is_gt': self.pos_is_gt,
            'num_gts': self.num_gts,
            'pos_assigned_gt_inds': self.pos_assigned_gt_inds,
            'num_pos': self.num_pos,
            'num_neg': self.num_neg,
            'avg_factor': self.avg_factor
        }

jittordet/models/task_utils/samplers/sampling_result.py
JittorCV/jittordet
from .initialize import (bias_init_with_prob, caffe2_xavier_init,
                         constant_init, kaiming_init, normal_init,
                         trunc_normal_init, uniform_init, xavier_init)
from .misc import (empty_instances, filter_scores_and_topk, images_to_levels,
                   multi_apply, select_single_mlvl, unmap, unpack_gt_instances)
from .nms import batched_nms, multiclass_nms
from .transforms import bbox2roi

__all__ = [
    'normal_init', 'constant_init', 'xavier_init', 'trunc_normal_init',
    'uniform_init', 'kaiming_init', 'caffe2_xavier_init',
    'bias_init_with_prob', 'unpack_gt_instances', 'empty_instances',
    'select_single_mlvl', 'filter_scores_and_topk', 'batched_nms',
    'multi_apply', 'unmap', 'images_to_levels', 'bbox2roi', 'multiclass_nms'
]

jittordet/models/utils/__init__.py
JittorCV/jittordet
# Modified from jdet/models/utils/weight_init.py
import jittor as jt
import numpy as np
from jittor import init, nn


def normal_init(module, mean=0, std=1, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.gauss_(module.weight, mean, std)
    if hasattr(module, 'bias') and isinstance(
            module.bias, jt.Var) and module.bias is not None:
        init.constant_(module.bias, bias)


def constant_init(module, val, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.constant_(module.weight, val)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def xavier_init(module, gain=1, bias=0, distribution='normal'):
    assert distribution in ['uniform', 'normal']
    if hasattr(module, 'weight') and module.weight is not None:
        if distribution == 'uniform':
            init.xavier_uniform_(module.weight, gain=gain)
        else:
            init.xavier_normal_(module.weight, gain=gain)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def trunc_normal_init(module: nn.Module,
                      mean: float = 0,
                      std: float = 1,
                      a: float = -2,
                      b: float = 2,
                      bias: float = 0) -> None:
    if hasattr(module, 'weight') and module.weight is not None:
        init.trunc_normal_(module.weight, mean, std, a, b)  # type: ignore
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)  # type: ignore


def uniform_init(module, a=0, b=1, bias=0):
    if hasattr(module, 'weight') and module.weight is not None:
        init.uniform_(module.weight, a, b)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def kaiming_init(module,
                 a=0,
                 mode='fan_out',
                 nonlinearity='relu',
                 bias=0,
                 distribution='normal'):
    assert distribution in ['uniform', 'normal']
    if hasattr(module, 'weight') and module.weight is not None:
        if distribution == 'uniform':
            init.kaiming_uniform_(
                module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
        else:
            init.kaiming_normal_(
                module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
    if hasattr(module, 'bias') and module.bias is not None:
        init.constant_(module.bias, bias)


def caffe2_xavier_init(module, bias=0):
    # `XavierFill` in Caffe2 corresponds to `kaiming_uniform_` in PyTorch
    # Acknowledgment to FAIR's internal code
    kaiming_init(
        module,
        a=1,
        mode='fan_in',
        nonlinearity='leaky_relu',
        bias=bias,
        distribution='uniform')


def bias_init_with_prob(prior_prob):
    """initialize conv/fc bias value according to a given probability value."""
    bias_init = float(-np.log((1 - prior_prob) / prior_prob))
    return bias_init

jittordet/models/utils/initialize.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/models/utils/misc.py
# Copyright (c) OpenMMLab. All rights reserved.
from functools import partial
from typing import List

import jittor as jt

from jittordet.structures import InstanceData, OptInstanceList, SampleList


def unpack_gt_instances(batch_data_samples: SampleList) -> tuple:
    """Unpack ``gt_instances``, ``gt_instances_ignore`` and ``img_metas`` based
    on ``batch_data_samples``

    Args:
        batch_data_samples (List[:obj:`DetDataSample`]): The Data
            Samples. It usually includes information such as
            `gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`.

    Returns:
        tuple:

            - batch_gt_instances (list[:obj:`InstanceData`]): Batch of
                gt_instance. It usually includes ``bboxes`` and ``labels``
                attributes.
            - batch_gt_instances_ignore (list[:obj:`InstanceData`]):
                Batch of gt_instances_ignore. It includes ``bboxes`` attribute
                data that is ignored during training and testing.
                Defaults to None.
            - batch_img_metas (list[dict]): Meta information of each image,
                e.g., image size, scaling factor, etc.
    """
    batch_gt_instances = []
    batch_gt_instances_ignore = []
    batch_img_metas = []
    for data_sample in batch_data_samples:
        batch_img_metas.append(data_sample.metainfo)
        batch_gt_instances.append(data_sample.gt_instances)
        if 'ignored_instances' in data_sample:
            batch_gt_instances_ignore.append(data_sample.ignored_instances)
        else:
            batch_gt_instances_ignore.append(None)

    return batch_gt_instances, batch_img_metas, batch_gt_instances_ignore


def select_single_mlvl(mlvl_tensors, batch_id, detach=True):
    """Extract a multi-scale single image tensor from a multi-scale batch
    tensor based on batch index.

    Note: The default value of detach is True, because the proposal gradient
    needs to be detached during the training of the two-stage model. E.g
    Cascade Mask R-CNN.

    Args:
        mlvl_tensors (list[Tensor]): Batch tensor for all scale levels,
           each is a 4D-tensor.
        batch_id (int): Batch index.
        detach (bool): Whether detach gradient. Default True.

    Returns:
        list[Tensor]: Multi-scale single image tensor.
    """
    assert isinstance(mlvl_tensors, (list, tuple))
    num_levels = len(mlvl_tensors)

    if detach:
        mlvl_tensor_list = [
            mlvl_tensors[i][batch_id].detach() for i in range(num_levels)
        ]
    else:
        mlvl_tensor_list = [
            mlvl_tensors[i][batch_id] for i in range(num_levels)
        ]
    return mlvl_tensor_list


def filter_scores_and_topk(scores, score_thr, topk, results=None):
    """Filter results using score threshold and topk candidates.

    Args:
        scores (Tensor): The scores, shape (num_bboxes, K).
        score_thr (float): The score filter threshold.
        topk (int): The number of topk candidates.
        results (dict or list or Tensor, Optional): The results to
           which the filtering rule is to be applied. The shape
           of each item is (num_bboxes, N).

    Returns:
        tuple: Filtered results

            - scores (Tensor): The scores after being filtered, \
                shape (num_bboxes_filtered, ).
            - labels (Tensor): The class labels, shape \
                (num_bboxes_filtered, ).
            - anchor_idxs (Tensor): The anchor indexes, shape \
                (num_bboxes_filtered, ).
            - filtered_results (dict or list or Tensor, Optional): \
                The filtered results. The shape of each item is \
                (num_bboxes_filtered, N).
    """
    valid_mask = scores > score_thr
    scores = scores[valid_mask]
    valid_idxs = jt.nonzero(valid_mask)

    num_topk = min(topk, valid_idxs.size(0))
    idxs, scores = jt.argsort(scores, descending=True)
    scores = scores[:num_topk]
    topk_idxs = valid_idxs[idxs[:num_topk]]
    keep_idxs, labels = topk_idxs.unbind(dim=1)

    filtered_results = None
    if results is not None:
        if isinstance(results, dict):
            filtered_results = {k: v[keep_idxs] for k, v in results.items()}
        elif isinstance(results, list):
            filtered_results = [result[keep_idxs] for result in results]
        elif isinstance(results, jt.Var):
            filtered_results = results[keep_idxs]
        else:
            raise NotImplementedError(f'Only supports dict or list or Tensor, '
                                      f'but get {type(results)}.')
    return scores, labels, keep_idxs, filtered_results


def multi_apply(func, *args, **kwargs):
    """Apply function to a list of arguments.

    Note:
        This function applies the ``func`` to multiple inputs and
        map the multiple outputs of the ``func`` into different
        list. Each list contains the same type of outputs corresponding
        to different inputs.

    Args:
        func (Function): A function that will be applied to a list of
            arguments

    Returns:
        tuple(list): A tuple containing multiple list, each list contains \
            a kind of returned results by the function
    """
    pfunc = partial(func, **kwargs) if kwargs else func
    map_results = map(pfunc, *args)
    return tuple(map(list, zip(*map_results)))


def unmap(data, count, inds, fill=0):
    """Unmap a subset of item (data) back to the original set of items (of size
    count)"""
    if data.ndim == 1:
        ret = jt.full((count, ), fill, dtype=data.dtype)
        ret[inds.astype(jt.bool)] = data
    else:
        new_size = (count, ) + data.size()[1:]
        ret = jt.full(new_size, fill, dtype=data.dtype)
        ret[inds.astype(jt.bool), :] = data
    return ret


def images_to_levels(target, num_levels):
    """Convert targets by image to targets by feature level.

    [target_img0, target_img1] -> [target_level0, target_level1, ...]
    """
    target = jt.stack(target, dim=0)
    level_targets = []
    start = 0
    for n in num_levels:
        end = start + n
        # level_targets.append(target[:, start:end].squeeze(0))
        level_targets.append(target[:, start:end])
        start = end
    return level_targets


def empty_instances(batch_img_metas: List[dict],
                    instance_results: OptInstanceList = None,
                    num_classes: int = 80,
                    score_per_cls: bool = False) -> List[InstanceData]:
    """Handle predicted instances when RoI is empty.

    Note: If ``instance_results`` is not None, it will be modified
    in place internally, and then return ``instance_results``
    Args:
        batch_img_metas (list[dict]): List of image information.
        device (torch.device): Device of tensor.
        task_type (str): Expected returned task type. it currently
            supports bbox and mask.
        instance_results (list[:obj:`InstanceData`]): List of instance
            results.
        mask_thr_binary (int, float): mask binarization threshold.
            Defaults to 0.
        box_type (str or type): The empty box type. Defaults to `hbox`.
        use_box_type (bool): Whether to warp boxes with the box type.
            Defaults to False.
        num_classes (int): num_classes of bbox_head. Defaults to 80.
        score_per_cls (bool):  Whether to generate classwise score for
            the empty instance. ``score_per_cls`` will be True when the model
            needs to produce raw results without nms. Defaults to False.
    Returns:
        list[:obj:`InstanceData`]: Detection results of each image
    """
    if instance_results is not None:
        assert len(instance_results) == len(batch_img_metas)

    results_list = []
    for img_id in range(len(batch_img_metas)):
        if instance_results is not None:
            results = instance_results[img_id]
            assert isinstance(results, InstanceData)
        else:
            results = InstanceData()
        results.bboxes = jt.zeros(0, 4)
        score_shape = (0, num_classes + 1) if score_per_cls else (0, )
        results.scores = jt.zeros(score_shape)
        results.labels = jt.zeros((0, ), dtype=jt.int64)
        results_list.append(results)
    return results_list

jittordet/models/utils/misc.py
JittorCV/jittordet
# Modified from OpenMMLab mmcv/ops/nms.py
from typing import Dict, Optional, Tuple

import jittor as jt

from jittordet.engine import ConfigType


def batched_nms(boxes: jt.Var,
                scores: jt.Var,
                idxs: jt.Var,
                nms_cfg: Optional[Dict],
                class_agnostic: bool = False) -> Tuple[jt.Var, jt.Var]:
    r"""Performs non-maximum suppression in a batched fashion.

    Modified from `torchvision/ops/boxes.py#L39
    <https://github.com/pytorch/vision/blob/
    505cd6957711af790211896d32b40291bea1bc21/torchvision/ops/boxes.py#L39>`_.
    In order to perform NMS independently per class, we add an offset to all
    the boxes. The offset is dependent only on the class idx, and is large
    enough so that boxes from different classes do not overlap.

    Note:
        In v1.4.1 and later, ``batched_nms`` supports skipping the NMS and
        returns sorted raw results when `nms_cfg` is None.

    Args:
        boxes (torch.Tensor): boxes in shape (N, 4) or (N, 5).
        scores (torch.Tensor): scores in shape (N, ).
        idxs (torch.Tensor): each index value correspond to a bbox cluster,
            and NMS will not be applied between elements of different idxs,
            shape (N, ).
        nms_cfg (dict | optional): Supports skipping the nms when `nms_cfg`
            is None, otherwise it should specify nms type and other
            parameters like `iou_thr`. Possible keys includes the following.

            - iou_threshold (float): IoU threshold used for NMS.
            - split_thr (float): threshold number of boxes. In some cases the
              number of boxes is large (e.g., 200k). To avoid OOM during
              training, the users could set `split_thr` to a small value.
              If the number of boxes is greater than the threshold, it will
              perform NMS on each group of boxes separately and sequentially.
              Defaults to 10000.
        class_agnostic (bool): if true, nms is class agnostic,
            i.e. IoU thresholding happens over all boxes,
            regardless of the predicted class. Defaults to False.

    Returns:
        tuple: kept dets and indice.

        - boxes (Tensor): Bboxes with score after nms, has shape
          (num_bboxes, 5). last dimension 5 arrange as
          (x1, y1, x2, y2, score)
        - keep (Tensor): The indices of remaining boxes in input
          boxes.
    """
    # skip nms when nms_cfg is None
    if nms_cfg is None:
        scores, inds = scores.sort(descending=True)
        boxes = boxes[inds]
        return boxes, scores, inds

    nms_cfg_ = nms_cfg.copy()
    class_agnostic = nms_cfg_.pop('class_agnostic', class_agnostic)
    if class_agnostic:
        boxes_for_nms = boxes
    else:
        # When using rotated boxes, only apply offsets on center.
        if boxes.size(-1) == 5:
            # Strictly, the maximum coordinates of the rotating box
            # (x,y,w,h,a) should be calculated by polygon coordinates.
            # But the conversion from rotated box to polygon will
            # slow down the speed.
            # So we use max(x,y) + max(w,h) as max coordinate
            # which is larger than polygon max coordinate
            # max(x1, y1, x2, y2,x3, y3, x4, y4)
            max_coordinate = boxes[..., :2].max() + boxes[..., 2:4].max()
            offsets = idxs.astype(boxes.dtype) * (max_coordinate + 1)
            boxes_ctr_for_nms = boxes[..., :2] + offsets[:, None]
            boxes_for_nms = jt.concat([boxes_ctr_for_nms, boxes[..., 2:5]],
                                      dim=-1)
        else:
            max_coordinate = boxes.max()
            offsets = idxs.astype(boxes.dtype) * (max_coordinate + 1)
            boxes_for_nms = boxes + offsets[:, None]

    # only support jt.nms for now
    nms_cfg_.pop('type', 'nms')
    dets = jt.concat([boxes_for_nms, scores[:, None]], dim=1)
    keep = jt.nms(dets, **nms_cfg_)
    boxes = boxes[keep]
    scores = scores[keep]
    return boxes, scores, keep


def multiclass_nms(multi_bboxes: jt.Var,
                   multi_scores: jt.Var,
                   score_thr: float,
                   nms_cfg: ConfigType,
                   max_num: int = -1,
                   score_factors: Optional[jt.Var] = None,
                   return_inds: bool = False):
    """NMS for multi-class bboxes.

    Args:
        multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
        multi_scores (Tensor): shape (n, #class), where the last column
            contains scores of the background class, but this will be ignored.
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_cfg (Union[:obj:`ConfigDict`, dict]): a dict that contains
            the arguments of nms operations.
        max_num (int, optional): if there are more than max_num bboxes after
            NMS, only top max_num will be kept. Default to -1.
        score_factors (Tensor, optional): The factors multiplied to scores
            before applying NMS. Default to None.
        return_inds (bool, optional): Whether return the indices of kept
            bboxes. Default to False.
        box_dim (int): The dimension of boxes. Defaults to 4.

    Returns:
        Union[Tuple[Tensor, Tensor, Tensor], Tuple[Tensor, Tensor]]:
            (dets, labels, indices (optional)), tensors of shape (k, 5),
            (k), and (k). Dets are boxes with scores. Labels are 0-based.
    """
    num_classes = multi_scores.size(1) - 1
    # exclude background category
    if multi_bboxes.shape[1] > 4:
        bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
    else:
        bboxes = multi_bboxes[:, None].expand(
            multi_scores.size(0), num_classes, 4)

    scores = multi_scores[:, :-1]

    labels = jt.arange(num_classes, dtype=jt.int64)
    labels = labels.view(1, -1).expand_as(scores)

    bboxes = bboxes.reshape(-1, 4)
    scores = scores.reshape(-1)
    labels = labels.reshape(-1)
    valid_mask = scores > score_thr
    # multiply score_factor after threshold to preserve more bboxes, improve
    # mAP by 1% for YOLOv3
    if score_factors is not None:
        # expand the shape to match original shape of score
        score_factors = score_factors.view(-1, 1).expand(
            multi_scores.size(0), num_classes)
        score_factors = score_factors.reshape(-1)
        scores = scores * score_factors

    # NonZero not supported  in TensorRT
    inds = valid_mask.nonzero().squeeze(1)
    bboxes, scores, labels = bboxes[inds], scores[inds], labels[inds]

    if bboxes.numel() == 0:
        if return_inds:
            return bboxes, scores, labels, inds
        else:
            return bboxes, scores, labels

    bboxes, scores, keep = batched_nms(bboxes, scores, labels, nms_cfg)

    if max_num > 0:
        bboxes = bboxes[:max_num]
        scores = scores[:max_num]
        keep = keep[:max_num]

    if return_inds:
        return bboxes, scores, labels[keep], inds[keep]
    else:
        return bboxes, scores, labels[keep]

jittordet/models/utils/nms.py
JittorCV/jittordet
from typing import List

import jittor as jt


def bbox2roi(bbox_list: List[jt.Var]) -> jt.Var:
    """Convert a list of bboxes to roi format.

    Args:
        bbox_list (List[Union[Tensor, :obj:`BaseBoxes`]): a list of bboxes
            corresponding to a batch of images.

    Returns:
        Tensor: shape (n, box_dim + 1), where ``box_dim`` depends on the
        different box types. For example, If the box type in ``bbox_list``
        is HorizontalBoxes, the output shape is (n, 5). Each row of data
        indicates [batch_ind, x1, y1, x2, y2].
    """
    rois_list = []
    for img_id, bboxes in enumerate(bbox_list):
        img_inds = jt.full((bboxes.size(0), 1), img_id, dtype=bboxes.dtype)
        rois = jt.concat([img_inds, bboxes], dim=-1)
        rois_list.append(rois)
    rois = jt.concat(rois_list, 0)
    return rois

jittordet/models/utils/transforms.py
JittorCV/jittordet
from .roi_align import ROIAlign, roi_align
from .roi_pool import ROIPool, roi_pool

__all__ = ['ROIAlign', 'roi_align', 'ROIPool', 'roi_pool']

jittordet/ops/__init__.py
JittorCV/jittordet
import numpy as np
import shapely.geometry as shgeo
from jdet.ops.bbox_transforms import *


def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
    assert mode in ['iou', 'iof']
    assert get_bbox_type(bboxes1) != 'notype'
    assert get_bbox_type(bboxes2) != 'notype'
    rows = bboxes1.shape[0]
    cols = bboxes2.shape[0]
    if is_aligned:
        assert rows == cols

    if rows * cols == 0:
        return np.zeros((rows, 1), dtype=np.float32) \
                if is_aligned else np.zeros((rows, cols), dtype=np.float32)

    hbboxes1 = bbox2type(bboxes1, 'hbb')
    hbboxes2 = bbox2type(bboxes2, 'hbb')
    if not is_aligned:
        hbboxes1 = hbboxes1[:, None, :]
    lt = np.maximum(hbboxes1[..., :2], hbboxes2[..., :2])
    rb = np.minimum(hbboxes1[..., 2:], hbboxes2[..., 2:])
    wh = np.clip(rb - lt, 0, np.inf)
    h_overlaps = wh[..., 0] * wh[..., 1]

    if get_bbox_type(bboxes1) == 'hbb' and get_bbox_type(bboxes2) == 'hbb':
        overlaps = h_overlaps
        areas1 = (hbboxes1[..., 2] - hbboxes1[..., 0]) * (
            hbboxes1[..., 3] - hbboxes1[..., 1])

        if mode == 'iou':
            areas2 = (hbboxes2[..., 2] - hbboxes2[..., 0]) * (
                hbboxes2[..., 3] - hbboxes2[..., 1])
            unions = areas1 + areas2 - overlaps
        else:
            unions = areas1

    else:
        polys1 = bbox2type(bboxes1, 'poly')
        polys2 = bbox2type(bboxes2, 'poly')
        sg_polys1 = [shgeo.Polygon(p) for p in polys1.reshape(rows, -1, 2)]
        sg_polys2 = [shgeo.Polygon(p) for p in polys2.reshape(cols, -1, 2)]

        overlaps = np.zeros(h_overlaps.shape)
        for p in zip(*np.nonzero(h_overlaps)):
            overlaps[p] = sg_polys1[p[0]].intersection(sg_polys2[p[-1]]).area

        if mode == 'iou':
            unions = np.zeros(h_overlaps.shape, dtype=np.float32)
            for p in zip(*np.nonzero(h_overlaps)):
                unions[p] = sg_polys1[p[0]].union(sg_polys2[p[-1]]).area
        else:
            unions = np.array([p.area for p in sg_polys1], dtype=np.float32)
            if not is_aligned:
                unions = unions[..., None]

    unions = np.clip(unions, eps, np.inf)
    outputs = overlaps / unions
    if outputs.ndim == 1:
        outputs = outputs[..., None]
    return outputs


def bbox_areas(bboxes):
    bbox_type = get_bbox_type(bboxes)
    assert bbox_type != 'notype'

    if bbox_type == 'hbb':
        areas = (bboxes[..., 2] - bboxes[..., 0]) * (
            bboxes[..., 3] - bboxes[..., 1])

    if bbox_type == 'obb':
        areas = bboxes[..., 2] * bboxes[..., 3]

    if bbox_type == 'poly':
        areas = np.zeros(bboxes.shape[:-1], dtype=np.float32)
        bboxes = bboxes.reshape(*bboxes.shape[:-1], 4, 2)
        for i in range(4):
            areas += 0.5 * (
                bboxes[..., i, 0] * bboxes[..., (i + 1) % 4, 1] -
                bboxes[..., (i + 1) % 4, 0] * bboxes[..., i, 1])
        areas = np.abs(areas)
    return areas


def bbox_nms(bboxes, scores, iou_thr=0.5, score_thr=0.01):
    assert get_bbox_type(bboxes) != 'notype'
    order = scores.argsort()[::-1]
    order = order[scores[order] > score_thr]
    keep = []

    while order.size > 0:
        i = order[0]
        keep.append(i)

        keep_bbox = bboxes[[i]]
        other_bboxes = bboxes[order[1:]]
        ious = bbox_overlaps(keep_bbox, other_bboxes)

        idx = np.where(ious <= iou_thr)[1]
        order = order[idx + 1]

    return np.array(keep, dtype=np.int64)


def bbox_area_nms(bboxes, iou_thr=0.5):
    assert get_bbox_type(bboxes) != 'notype'
    areas = bbox_areas(bboxes)
    order = areas.argsort()[::-1]
    keep = []

    while order.size > 0:
        i = order[0]
        keep.append(i)

        keep_bbox = bboxes[[i]]
        other_bboxes = bboxes[order[1:]]
        ious = bbox_overlaps(keep_bbox, other_bboxes)

        idx = np.where(ious <= iou_thr)[1]
        order = order[idx + 1]

    return np.array(keep, dtype=np.int64)

jittordet/ops/bbox_geometry.py
JittorCV/jittordet
import copy  # TODO: remove copy(?)
import math
from typing import Tuple

import cv2
import jittor as jt
import numpy as np  # TODO: remove numpy


def dbbox2delta_v3(proposals, gt, means=[0, 0, 0, 0, 0], stds=[1, 1, 1, 1, 1]):
    proposals = proposals.float()
    gt = gt.float()
    gt_widths = gt[..., 2]
    gt_heights = gt[..., 3]
    gt_angle = gt[..., 4]

    proposals_widths = proposals[..., 2]
    proposals_heights = proposals[..., 3]
    proposals_angle = proposals[..., 4]

    coord = gt[..., 0:2] - proposals[..., 0:2]
    dx = (jt.cos(proposals[..., 4]) * coord[..., 0] +
          jt.sin(proposals[..., 4]) * coord[..., 1]) / proposals_widths
    dy = (-jt.sin(proposals[..., 4]) * coord[..., 0] +
          jt.cos(proposals[..., 4]) * coord[..., 1]) / proposals_heights
    dw = jt.log(gt_widths / proposals_widths)
    dh = jt.log(gt_heights / proposals_heights)
    dangle = gt_angle - proposals_angle
    deltas = jt.stack((dx, dy, dw, dh, dangle), -1)

    means = jt.array(means).unsqueeze(0)
    stds = jt.array(stds).unsqueeze(0)
    deltas = (deltas - means) / stds

    return deltas


def hbb2obb_v2(boxes):
    num_boxes = boxes.size(0)
    ex_heights = boxes[..., 2] - boxes[..., 0] + 1.0
    ex_widths = boxes[..., 3] - boxes[..., 1] + 1.0
    ex_ctr_x = boxes[..., 0] + 0.5 * (ex_heights - 1.0)
    ex_ctr_y = boxes[..., 1] + 0.5 * (ex_widths - 1.0)
    c_bboxes = jt.contrib.concat(
        (ex_ctr_x.unsqueeze(1), ex_ctr_y.unsqueeze(1), ex_widths.unsqueeze(1),
         ex_heights.unsqueeze(1)), 1)
    initial_angles = -jt.ones((num_boxes, 1)) * np.pi / 2
    dbboxes = jt.contrib.concat((c_bboxes, initial_angles), 1)

    return dbboxes


def mask2poly_single(binary_mask):
    res = cv2.findContours(
        binary_mask.copy(), cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_NONE)  #TODO: check binary_mask.copy()
    if (len(res) == 2):
        contours, hierarchy = res
    else:
        _, contours, hierarchy = res
    max_contour = max(contours, key=len)
    rect = cv2.minAreaRect(max_contour)
    poly = cv2.boxPoints(rect)
    return poly


def mask2poly(binary_mask_list):
    polys = map(mask2poly_single, binary_mask_list)
    return list(polys)


def obb2poly_single(gt_obb):
    '''
    Args:
        gt_obb (np.array): [x, y, h, w, a]
    Rets:
        poly (np.array): shape is [4, 2]
    '''
    center = np.array([gt_obb[0], gt_obb[1]])
    vx = np.array([(gt_obb[2] - 1) * math.cos(gt_obb[4]) / 2,
                   (gt_obb[2] - 1) * math.sin(gt_obb[4]) / 2])
    vy = np.array([(gt_obb[3] - 1) * math.sin(gt_obb[4]) / 2,
                   -(gt_obb[3] - 1) * math.cos(gt_obb[4]) / 2])
    p1 = center + vx + vy
    p2 = center + vx - vy
    p3 = center - vx - vy
    p4 = center - vx + vy
    return np.stack((p1, p2, p3, p4))


def obb2poly_v0(gt_obb_list):
    polys = map(obb2poly_single, gt_obb_list)
    return list(polys)


def polygonToRotRectangle_batch(bbox, with_module=True):
    bbox = np.array(bbox, dtype=np.float32)
    bbox = np.reshape(bbox, newshape=(-1, 2, 4), order='F')
    # angle = math.atan2(-(bbox[0,1]-bbox[0,0]),bbox[1,1]-bbox[1,0])
    # print('bbox: ', bbox)
    angle = np.arctan2(-(bbox[:, 0, 1] - bbox[:, 0, 0]),
                       bbox[:, 1, 1] - bbox[:, 1, 0])
    # angle = np.arctan2(-(bbox[:, 0,1]-bbox[:, 0,0]),bbox[:, 1,1]-bbox[:, 1,0])
    # center = [[0],[0]] ## shape [2, 1]
    # print('angle: ', angle)
    center = np.zeros((bbox.shape[0], 2, 1))
    for i in range(4):
        center[:, 0, 0] += bbox[:, 0, i]
        center[:, 1, 0] += bbox[:, 1, i]

    center = np.array(center, dtype=np.float32) / 4.0

    # R = np.array([[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]], dtype=np.float32)
    R = np.array([[np.cos(angle), -np.sin(angle)],
                  [np.sin(angle), np.cos(angle)]],
                 dtype=np.float32)

    normalized = np.matmul(R.transpose((2, 1, 0)), bbox - center)

    xmin = np.min(normalized[:, 0, :], axis=1)
    # print('diff: ', (xmin - normalized[:, 0, 3]))
    # assert sum((abs(xmin - normalized[:, 0, 3])) > eps) == 0
    xmax = np.max(normalized[:, 0, :], axis=1)
    # assert sum(abs(xmax - normalized[:, 0, 1]) > eps) == 0
    # print('diff2: ', xmax - normalized[:, 0, 1])
    ymin = np.min(normalized[:, 1, :], axis=1)
    # assert sum(abs(ymin - normalized[:, 1, 3]) > eps) == 0
    # print('diff3: ', ymin - normalized[:, 1, 3])
    ymax = np.max(normalized[:, 1, :], axis=1)
    # assert sum(abs(ymax - normalized[:, 1, 1]) > eps) == 0
    # print('diff4: ', ymax - normalized[:, 1, 1])

    w = xmax - xmin + 1
    h = ymax - ymin + 1

    w = w[:, np.newaxis]
    h = h[:, np.newaxis]
    # TODO: check it
    if with_module:
        angle = angle[:, np.newaxis] % (2 * np.pi)
    else:
        angle = angle[:, np.newaxis]
    dboxes = np.concatenate((center[:, 0].astype(
        np.float), center[:, 1].astype(np.float), w, h, angle),
                            axis=1)
    return dboxes


def cal_line_length(point1, point2):
    return math.sqrt(
        math.pow(point1[0] - point2[0], 2) +
        math.pow(point1[1] - point2[1], 2))


def get_best_begin_point_single(coordinate):
    x1 = coordinate[0][0]
    y1 = coordinate[0][1]
    x2 = coordinate[1][0]
    y2 = coordinate[1][1]
    x3 = coordinate[2][0]
    y3 = coordinate[2][1]
    x4 = coordinate[3][0]
    y4 = coordinate[3][1]
    xmin = min(x1, x2, x3, x4)
    ymin = min(y1, y2, y3, y4)
    xmax = max(x1, x2, x3, x4)
    ymax = max(y1, y2, y3, y4)
    combinate = [[[x1, y1], [x2, y2], [x3, y3], [x4, y4]],
                 [[x2, y2], [x3, y3], [x4, y4], [x1, y1]],
                 [[x3, y3], [x4, y4], [x1, y1], [x2, y2]],
                 [[x4, y4], [x1, y1], [x2, y2], [x3, y3]]]
    dst_coordinate = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
    force = 100000000.0
    force_flag = 0
    for i in range(4):
        temp_force = cal_line_length(
            combinate[i][0], dst_coordinate[0]) + cal_line_length(
                combinate[i][1], dst_coordinate[1]) + cal_line_length(
                    combinate[i][2], dst_coordinate[2]) + cal_line_length(
                        combinate[i][3], dst_coordinate[3])
        if temp_force < force:
            force = temp_force
            force_flag = i
    if force_flag != 0:
        pass
        # print("choose one direction!")
    return combinate[force_flag]


def TuplePoly2Poly(poly):
    outpoly = [
        poly[0][0], poly[0][1], poly[1][0], poly[1][1], poly[2][0], poly[2][1],
        poly[3][0], poly[3][1]
    ]
    return outpoly


def get_best_begin_point_warp_single(coordinate):
    return TuplePoly2Poly(get_best_begin_point_single(coordinate))


def get_best_begin_point(coordinate_list):
    best_coordinate_list = map(get_best_begin_point_warp_single,
                               coordinate_list)
    best_coordinate_list = np.stack(list(best_coordinate_list))
    return best_coordinate_list


def bbox2delta(proposals, gt, means=[0, 0, 0, 0], stds=[1, 1, 1, 1]):
    assert proposals.size() == gt.size()

    proposals = proposals.float()
    gt = gt.float()
    px = (proposals[..., 0] + proposals[..., 2]) * 0.5
    py = (proposals[..., 1] + proposals[..., 3]) * 0.5
    pw = proposals[..., 2] - proposals[..., 0] + 1.0
    ph = proposals[..., 3] - proposals[..., 1] + 1.0

    gx = (gt[..., 0] + gt[..., 2]) * 0.5
    gy = (gt[..., 1] + gt[..., 3]) * 0.5
    gw = gt[..., 2] - gt[..., 0] + 1.0
    gh = gt[..., 3] - gt[..., 1] + 1.0

    dx = (gx - px) / pw
    dy = (gy - py) / ph
    dw = jt.log(gw / pw)
    dh = jt.log(gh / ph)
    deltas = jt.stack([dx, dy, dw, dh], dim=-1)

    means = jt.float32(means).unsqueeze(0)
    stds = jt.float32(stds).unsqueeze(0)
    deltas = deltas.subtract(means).divide(stds)

    return deltas


def dbbox2delta_v2(proposals, gt, means=[0, 0, 0, 0, 0], stds=[1, 1, 1, 1, 1]):
    gt_widths = gt[..., 2]
    gt_heights = gt[..., 3]
    gt_angle = gt[..., 4]

    roi_widths = proposals[..., 2]
    roi_heights = proposals[..., 3]
    roi_angle = proposals[..., 4]

    coord = gt[..., 0:2] - proposals[..., 0:2]
    targets_dx = (jt.cos(roi_angle) * coord[..., 0] +
                  jt.sin(roi_angle) * coord[:, 1]) / roi_widths
    targets_dy = (-jt.sin(roi_angle) * coord[..., 0] +
                  jt.cos(roi_angle) * coord[:, 1]) / roi_heights
    targets_dw = jt.log(gt_widths / roi_widths)
    targets_dh = jt.log(gt_heights / roi_heights)
    targets_dangle = (gt_angle - roi_angle)
    dist = targets_dangle % (2 * np.pi)
    dist = jt.minimum(dist, np.pi * 2 - dist)

    inds = jt.sin(targets_dangle) < 0
    dist[inds] = -dist[inds]
    # TODO: change the norm value
    dist = dist / (np.pi / 2.)
    deltas = jt.stack((targets_dx, targets_dy, targets_dw, targets_dh, dist),
                      -1)

    means = jt.array(means).unsqueeze(0)
    stds = jt.array(stds).unsqueeze(0)
    deltas = (deltas - means) / stds

    return deltas


def choose_best_match_batch(Rrois, gt_rois):
    Rroi_angles = Rrois[:, 4].unsqueeze(1)

    gt_xs, gt_ys, gt_ws, gt_hs, gt_angles = gt_rois[:, 0].copy(
    ), gt_rois[:, 1].copy(), gt_rois[:, 2].copy(), gt_rois[:, 3].copy(
    ), gt_rois[:, 4].copy()

    gt_angle_extent = jt.contrib.concat(
        (gt_angles[:, np.newaxis], (gt_angles + np.pi / 2.)[:, np.newaxis],
         (gt_angles + np.pi)[:, np.newaxis],
         (gt_angles + np.pi * 3 / 2.)[:, np.newaxis]), 1)
    dist = (Rroi_angles - gt_angle_extent) % (2 * np.pi)
    dist = jt.minimum(dist, np.pi * 2 - dist)
    min_index = jt.argmin(dist, 1)[0]

    gt_rois_extent0 = copy.deepcopy(gt_rois)
    gt_rois_extent1 = jt.contrib.concat((gt_xs.unsqueeze(1), gt_ys.unsqueeze(1), \
                                 gt_hs.unsqueeze(1), gt_ws.unsqueeze(1), gt_angles.unsqueeze(1) + np.pi/2.), 1)
    gt_rois_extent2 = jt.contrib.concat((gt_xs.unsqueeze(1), gt_ys.unsqueeze(1), \
                                 gt_ws.unsqueeze(1), gt_hs.unsqueeze(1), gt_angles.unsqueeze(1) + np.pi), 1)
    gt_rois_extent3 = jt.contrib.concat((gt_xs.unsqueeze(1), gt_ys.unsqueeze(1), \
                                 gt_hs.unsqueeze(1), gt_ws.unsqueeze(1), gt_angles.unsqueeze(1) + np.pi * 3/2.), 1)
    gt_rois_extent = jt.contrib.concat(
        (gt_rois_extent0.unsqueeze(1), gt_rois_extent1.unsqueeze(1),
         gt_rois_extent2.unsqueeze(1), gt_rois_extent3.unsqueeze(1)), 1)

    gt_rois_new = jt.zeros_like(gt_rois)
    for curiter, index in enumerate(min_index):
        gt_rois_new[curiter, :] = gt_rois_extent[curiter, index.item(), :]

    gt_rois_new[:, 4] = gt_rois_new[:, 4] % (2 * np.pi)

    return gt_rois_new


def best_match_dbbox2delta(Rrois,
                           gt,
                           means=[0, 0, 0, 0, 0],
                           stds=[1, 1, 1, 1, 1]):
    gt_boxes_new = choose_best_match_batch(Rrois, gt)
    bbox_targets = dbbox2delta_v2(Rrois, gt_boxes_new, means, stds)

    return bbox_targets


def dbbox2result(dbboxes, labels, num_classes):
    dbboxes = dbboxes
    labels = labels
    return dbboxes[:, :8], dbboxes[:, -1].flatten(), labels


def delta2dbbox_v3(Rrois,
                   deltas,
                   means=[0, 0, 0, 0, 0],
                   stds=[1, 1, 1, 1, 1],
                   max_shape=None,
                   wh_ratio_clip=16 / 1000):
    means = jt.array(means, dtype=deltas.dtype).repeat(1, deltas.size(1) // 5)
    stds = jt.array(stds, dtype=deltas.dtype).repeat(1, deltas.size(1) // 5)
    denorm_deltas = deltas * stds + means

    dx = denorm_deltas[:, 0::5]
    dy = denorm_deltas[:, 1::5]
    dw = denorm_deltas[:, 2::5]
    dh = denorm_deltas[:, 3::5]
    dangle = denorm_deltas[:, 4::5]

    max_ratio = np.abs(np.log(wh_ratio_clip))
    dw = dw.clamp(min_v=-max_ratio, max_v=max_ratio)
    dh = dh.clamp(min_v=-max_ratio, max_v=max_ratio)
    Rroi_x = (Rrois[:, 0]).unsqueeze(1).expand_as(dx)
    Rroi_y = (Rrois[:, 1]).unsqueeze(1).expand_as(dy)
    Rroi_w = (Rrois[:, 2]).unsqueeze(1).expand_as(dw)
    Rroi_h = (Rrois[:, 3]).unsqueeze(1).expand_as(dh)
    Rroi_angle = (Rrois[:, 4]).unsqueeze(1).expand_as(dangle)
    # import pdb
    # pdb.set_trace()
    gx = dx * Rroi_w * jt.cos(Rroi_angle) \
         - dy * Rroi_h * jt.sin(Rroi_angle) + Rroi_x
    gy = dx * Rroi_w * jt.sin(Rroi_angle) \
         + dy * Rroi_h * jt.cos(Rroi_angle) + Rroi_y
    gw = Rroi_w * dw.exp()
    gh = Rroi_h * dh.exp()

    # TODO: check the hard code
    # gangle = (2 * np.pi) * dangle + Rroi_angle
    gangle = dangle + Rroi_angle
    # gangle = gangle % ( 2 * np.pi)

    if max_shape is not None:
        pass

    bboxes = jt.stack([gx, gy, gw, gh, gangle], dim=-1).view_as(deltas)
    return bboxes


def delta2dbbox_v2(Rrois,
                   deltas,
                   means=[0, 0, 0, 0, 0],
                   stds=[1, 1, 1, 1, 1],
                   max_shape=None,
                   wh_ratio_clip=16 / 1000):
    means = jt.array(means, dtype=deltas.dtype).repeat(1, deltas.size(1) // 5)
    stds = jt.array(stds, dtype=deltas.dtype).repeat(1, deltas.size(1) // 5)
    denorm_deltas = deltas * stds + means

    dx = denorm_deltas[:, 0::5]
    dy = denorm_deltas[:, 1::5]
    dw = denorm_deltas[:, 2::5]
    dh = denorm_deltas[:, 3::5]
    dangle = denorm_deltas[:, 4::5]

    max_ratio = np.abs(np.log(wh_ratio_clip))
    dw = dw.clamp(min_v=-max_ratio, max_v=max_ratio)
    dh = dh.clamp(min_v=-max_ratio, max_v=max_ratio)
    Rroi_x = (Rrois[:, 0]).unsqueeze(1).expand_as(dx)
    Rroi_y = (Rrois[:, 1]).unsqueeze(1).expand_as(dy)
    Rroi_w = (Rrois[:, 2]).unsqueeze(1).expand_as(dw)
    Rroi_h = (Rrois[:, 3]).unsqueeze(1).expand_as(dh)
    Rroi_angle = (Rrois[:, 4]).unsqueeze(1).expand_as(dangle)
    gx = dx * Rroi_w * jt.cos(Rroi_angle) \
         - dy * Rroi_h * jt.sin(Rroi_angle) + Rroi_x
    gy = dx * Rroi_w * jt.sin(Rroi_angle) \
         + dy * Rroi_h * jt.cos(Rroi_angle) + Rroi_y
    gw = Rroi_w * dw.exp()
    gh = Rroi_h * dh.exp()

    gangle = (np.pi / 2.) * dangle + Rroi_angle

    if max_shape is not None:
        pass

    bboxes = jt.stack([gx, gy, gw, gh, gangle], dim=-1).view_as(deltas)
    return bboxes


def delta2bbox(rois,
               deltas,
               means=[0, 0, 0, 0],
               stds=[1, 1, 1, 1],
               max_shape=None,
               wh_ratio_clip=16 / 1000):
    means = jt.array(means).repeat(1, deltas.size(1) // 4)
    stds = jt.array(stds).repeat(1, deltas.size(1) // 4)
    denorm_deltas = deltas * stds + means
    dx = denorm_deltas[:, 0::4]
    dy = denorm_deltas[:, 1::4]
    dw = denorm_deltas[:, 2::4]
    dh = denorm_deltas[:, 3::4]
    max_ratio = np.abs(np.log(wh_ratio_clip))
    dw = dw.clamp(min_v=-max_ratio, max_v=max_ratio)
    dh = dh.clamp(min_v=-max_ratio, max_v=max_ratio)
    px = ((rois[:, 0] + rois[:, 2]) * 0.5).unsqueeze(1).expand_as(dx)
    py = ((rois[:, 1] + rois[:, 3]) * 0.5).unsqueeze(1).expand_as(dy)
    pw = (rois[:, 2] - rois[:, 0] + 1.0).unsqueeze(1).expand_as(dw)
    ph = (rois[:, 3] - rois[:, 1] + 1.0).unsqueeze(1).expand_as(dh)
    gw = pw * dw.exp()
    gh = ph * dh.exp()
    gx = px + pw * dx  # gx = px + pw * dx
    gy = py + ph * dy  # gy = py + ph * dy
    x1 = gx - gw * 0.5 + 0.5
    y1 = gy - gh * 0.5 + 0.5
    x2 = gx + gw * 0.5 - 0.5
    y2 = gy + gh * 0.5 - 0.5
    if max_shape is not None:
        x1 = x1.clamp(min_v=0, max_v=max_shape[1] - 1)
        y1 = y1.clamp(min_v=0, max_v=max_shape[0] - 1)
        x2 = x2.clamp(min_v=0, max_v=max_shape[1] - 1)
        y2 = y2.clamp(min_v=0, max_v=max_shape[0] - 1)
    bboxes = jt.stack([x1, y1, x2, y2], dim=-1).view_as(deltas)
    return bboxes


def bbox2roi(bbox_list):
    """Convert a list of bboxes to roi format.

    Args:
        bbox_list (list[Tensor]): a list of bboxes corresponding to a batch
            of images.

    Returns:
        Tensor: shape (n, 5), [batch_ind, x1, y1, x2, y2]
    """
    rois_list = []
    for img_id, bboxes in enumerate(bbox_list):
        if bboxes.size(0) > 0:
            img_inds = jt.ones(
                (bboxes.size(0), 1), dtype=bboxes.dtype) * img_id
            rois = jt.contrib.concat([img_inds, bboxes[:, :4]], dim=-1)
        else:
            rois = bboxes.new_zeros((0, 5))
        rois_list.append(rois)
    rois = jt.contrib.concat(rois_list, 0)
    return rois


def gt_mask_bp_obbs(gt_masks, with_module=True):

    # trans gt_masks to gt_obbs
    gt_polys = mask2poly(gt_masks)
    gt_bp_polys = get_best_begin_point(gt_polys)
    gt_obbs = polygonToRotRectangle_batch(gt_bp_polys, with_module)

    return gt_obbs


def gt_mask_bp_obbs_list(gt_masks_list):

    gt_obbs_list = map(gt_mask_bp_obbs, gt_masks_list)

    return list(gt_obbs_list)


def roi2droi(rois):
    """
    :param rois: Tensor: shape (n, 5), [batch_ind, x1, y1, x2, y2]
    :return: drois: Tensor: shape (n, 6), [batch_ind, x, y, w, h, theta]
    """
    hbbs = rois[:, 1:]
    obbs = hbb2obb_v2(hbbs)

    return jt.contrib.concat((rois[:, 0].unsqueeze(1), obbs), 1)


def choose_best_Rroi_batch(Rroi):
    """There are many instances with large aspect ratio, so we choose the
    point, previous is long side, after is short side, so it makes sure h < w
    then angle % 180,

    :param Rroi: (x_ctr, y_ctr, w, h, angle)
            shape: (n, 5)
    :return: Rroi_new: Rroi with new representation
    """
    x_ctr, y_ctr, w, h, angle = copy.deepcopy(Rroi[:, 0]), copy.deepcopy(Rroi[:, 1]), \
                                copy.deepcopy(Rroi[:, 2]), copy.deepcopy(Rroi[:, 3]), copy.deepcopy(Rroi[:, 4])
    indexes = w < h

    Rroi[indexes, 2] = h[indexes]
    Rroi[indexes, 3] = w[indexes]
    Rroi[indexes, 4] = Rroi[indexes, 4] + np.pi / 2.
    # TODO: check the module
    Rroi[:, 4] = Rroi[:, 4] % np.pi

    return Rroi


def choose_best_obb_batch(ori_gt_obbs):
    """make the angle of obb close to -90.

    :param gt_obb: (x_ctr, y_ctr, w, h, angle)
            shape: (n, 5)
    """
    gt_obbs = copy.deepcopy(ori_gt_obbs)
    w, h = ori_gt_obbs[:, 2], ori_gt_obbs[:, 3]
    gt_obbs[:, 4] = (gt_obbs[:, 4] - np.pi * 1. / 4.) % np.pi
    indexes = gt_obbs[:, 4] >= np.pi / 2
    gt_obbs[indexes, 2] = h[indexes]
    gt_obbs[indexes, 3] = w[indexes]
    gt_obbs[indexes, 4] = gt_obbs[indexes, 4] - np.pi / 2.
    gt_obbs[:, 4] = gt_obbs[:, 4] - np.pi * 3. / 4.
    return gt_obbs


def dbbox2roi(dbbox_list):
    """Convert a list of dbboxes to droi format.

    :param dbbox_list: (list[Tensor]): a list of dbboxes corresponding to a batch of images
    :return: Tensor: shape (n, 6) [batch_ind, x_ctr, y_ctr, w, h, angle]
    """
    drois_list = []
    for img_id, dbboxes in enumerate(dbbox_list):
        if dbboxes.size(0) > 0:
            img_inds = jt.full((dbboxes.size(0), 1),
                               img_id,
                               dtype=dbboxes.dtype)
            drois = jt.contrib.concat([img_inds, dbboxes[:, :5]], dim=-1)
        else:
            drois = jt.zeros((0, 6), dtype=dbboxes.dtype)

        drois_list.append(drois)
    drois = jt.contrib.concat(drois_list, 0)
    return drois


def regular_theta(theta, mode='180', start=-np.pi / 2):
    assert mode in ['360', '180']
    cycle = 2 * np.pi if mode == '360' else np.pi

    theta = theta - start
    theta = theta % cycle
    return theta + start


def regular_obb(obboxes):
    x, y, w, h, theta = obboxes.unbind(dim=-1)

    w_regular = w * (w > h) + h * (1 - (w > h))
    h_regular = h * (w > h) + w * (1 - (w > h))
    theta_regular = theta * (w > h) + (theta + np.pi / 2) * (1 - (w > h))
    # w_regular = jt.where(w > h, w, h)
    # h_regular = jt.where(w > h, h, w)
    # theta_regular = jt.where(w > h, theta, theta + np.pi/2)
    theta_regular = regular_theta(theta_regular)
    return jt.stack([x, y, w_regular, h_regular, theta_regular], dim=-1)


def get_bbox_type(bboxes, with_score=False):

    dim = bboxes.size(-1)

    if with_score:
        dim -= 1
    if dim == 4:
        return 'hbb'
    if dim == 5:
        return 'obb'
    if dim == 8:
        return 'poly'
    return 'notype'


def get_bbox_dim(bbox_type, with_score=False):
    if bbox_type == 'hbb':
        dim = 4
    elif bbox_type == 'obb':
        dim = 5
    elif bbox_type == 'poly':
        dim = 8
    else:
        raise ValueError(f"don't know {bbox_type} bbox dim")

    if with_score:
        dim += 1
    return dim


def poly2obb(polys):

    polys_np = polys.numpy()

    order = polys_np.shape[:-1]
    num_points = polys_np.shape[-1] // 2
    polys_np = polys_np.reshape(-1, num_points, 2)
    polys_np = polys_np.astype(np.float32)

    obboxes = []
    for poly in polys_np:
        (x, y), (w, h), angle = cv2.minAreaRect(poly)
        if w >= h:
            angle = -angle
        else:
            w, h = h, w
            angle = -90 - angle
        theta = angle / 180 * np.pi
        obboxes.append([x, y, w, h, theta])

    if not obboxes:
        obboxes = np.zeros((0, 5))
    else:
        obboxes = np.array(obboxes)

    obboxes = obboxes.reshape(*order, 5)
    return jt.array(obboxes)


def rectpoly2obb(polys):

    theta = jt.arctan2(-(polys[..., 3] - polys[..., 1]),
                       polys[..., 2] - polys[..., 0])
    Cos, Sin = jt.cos(theta), jt.sin(theta)
    Matrix = jt.stack([Cos, -Sin, Sin, Cos], dim=-1)
    Matrix = Matrix.view(*Matrix.shape[:-1], 2, 2)

    x = polys[..., 0::2].mean(-1)
    y = polys[..., 1::2].mean(-1)
    center = jt.stack([x, y], dim=-1).unsqueeze(-2)
    center_polys = polys.view(*polys.shape[:-1], 4, 2) - center
    rotate_polys = jt.matmul(center_polys, Matrix.transpose(-1, -2))

    xmin = jt.min(rotate_polys[..., :, 0], dim=-1)
    xmax = jt.max(rotate_polys[..., :, 0], dim=-1)
    ymin = jt.min(rotate_polys[..., :, 1], dim=-1)
    ymax = jt.max(rotate_polys[..., :, 1], dim=-1)
    w = xmax - xmin
    h = ymax - ymin

    obboxes = jt.stack([x, y, w, h, theta], dim=-1)
    return regular_obb(obboxes)


def poly2hbb(polys):

    polys = polys.view(*polys.shape[:-1], polys.size(-1) // 2, 2)

    lt_point = jt.min(polys, dim=-2)
    rb_point = jt.max(polys, dim=-2)

    return jt.concat([lt_point, rb_point], dim=-1)


def obb2poly(obboxes):
    center, w, h, theta = jt.split(obboxes, [2, 1, 1, 1], dim=-1)
    Cos, Sin = jt.cos(theta), jt.sin(theta)

    vector1 = jt.concat([w / 2 * Cos, -w / 2 * Sin], dim=-1)
    vector2 = jt.concat([-h / 2 * Sin, -h / 2 * Cos], dim=-1)

    point1 = center + vector1 + vector2
    point2 = center + vector1 - vector2
    point3 = center - vector1 - vector2
    point4 = center - vector1 + vector2
    return jt.concat([point1, point2, point3, point4], dim=-1)


def obb2hbb(obboxes):
    center, w, h, theta = jt.split(obboxes, [2, 1, 1, 1], dim=-1)
    Cos, Sin = jt.cos(theta), jt.sin(theta)
    x_bias = jt.abs(w / 2 * Cos) + jt.abs(h / 2 * Sin)
    y_bias = jt.abs(w / 2 * Sin) + jt.abs(h / 2 * Cos)
    bias = jt.concat([x_bias, y_bias], dim=-1)
    return jt.concat([center - bias, center + bias], dim=-1)


def hbb2poly(hbboxes):
    l, t, r, b = hbboxes.unbind(-1)
    return jt.stack([l, t, r, t, r, b, l, b], dim=-1)


def hbb2obb(hbboxes):
    x = (hbboxes[..., 0] + hbboxes[..., 2]) * 0.5
    y = (hbboxes[..., 1] + hbboxes[..., 3]) * 0.5
    w = hbboxes[..., 2] - hbboxes[..., 0]
    h = hbboxes[..., 3] - hbboxes[..., 1]
    theta = jt.zeros_like(x)

    obboxes1 = jt.stack([x, y, w, h, theta], dim=-1)
    obboxes2 = jt.stack([x, y, h, w, theta - np.pi / 2], dim=-1)
    flag = (w >= h)[..., None]
    obboxes = flag * obboxes1 + (1 - flag) * obboxes2

    return obboxes


_type_func_map = {
    ('poly', 'obb'): poly2obb,
    ('poly', 'hbb'): poly2hbb,
    ('obb', 'poly'): obb2poly,
    ('obb', 'hbb'): obb2hbb,
    ('hbb', 'poly'): hbb2poly,
    ('hbb', 'obb'): hbb2obb
}


def bbox2type(bboxes, to_type):
    assert to_type in ['hbb', 'obb', 'poly']

    ori_type = get_bbox_type(bboxes)
    if ori_type == 'notype':
        raise ValueError('Not a bbox type')
    if ori_type == to_type:
        return bboxes
    trans_func = _type_func_map[(ori_type, to_type)]
    return trans_func(bboxes)


def get_bbox_areas(bboxes):
    btype = get_bbox_type(bboxes)
    if btype == 'hbb':
        wh = bboxes[..., 2:] - bboxes[..., :2]
        areas = wh[..., 0] * wh[..., 1]
    elif btype == 'obb':
        areas = bboxes[..., 2] * bboxes[..., 3]
    elif btype == 'poly':
        pts = bboxes.view(*bboxes.size()[:-1], 4, 2)
        roll_pts = torch.roll(pts, 1, dims=-2)
        xyxy = torch.sum(
            pts[..., 0] * roll_pts[..., 1] - roll_pts[..., 0] * pts[..., 1],
            dim=-1)
        areas = 0.5 * torch.abs(xyxy)
    else:
        raise ValueError('The type of bboxes is notype')

    return areas


def scale_boxes(boxes: jt.Var, scale_factor: Tuple[float, float]) -> jt.Var:
    repeat_num = int(boxes.size(-1) / 2)
    scale_factor = jt.Var(scale_factor).repeat((1, repeat_num))
    return boxes * scale_factor

jittordet/ops/bbox_transforms.py
JittorCV/jittordet
import jittor as jt

IOU_ROTATED_HEADER1 = r'''
// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
#undef out
#include <cassert>
#include <cmath>
#include <executor.h>
#include <algorithm>
#define CeilDIV(a,b) ((a+b-1)/b)
'''

IOU_ROTATED_HEADER2 = r'''
namespace {
int const threadsPerBlock = sizeof(unsigned long long) * 8;

template <typename T>
struct RotatedBox {
  T x_ctr, y_ctr, w, h, a;
};

template <typename T>
struct Point {
  T x, y;
  HOST_DEVICE_INLINE Point(const T& px = 0, const T& py = 0) : x(px), y(py) {}
  HOST_DEVICE_INLINE Point operator+(const Point& p) const {
    return Point(x + p.x, y + p.y);
  }
  HOST_DEVICE_INLINE Point& operator+=(const Point& p) {
    x += p.x;
    y += p.y;
    return *this;
  }
  HOST_DEVICE_INLINE Point operator-(const Point& p) const {
    return Point(x - p.x, y - p.y);
  }
  HOST_DEVICE_INLINE Point operator*(const T coeff) const {
    return Point(x * coeff, y * coeff);
  }
};

template <typename T>
HOST_DEVICE_INLINE T dot_2d(const Point<T>& A, const Point<T>& B) {
  return A.x * B.x + A.y * B.y;
}

template <typename T>
HOST_DEVICE_INLINE T cross_2d(const Point<T>& A, const Point<T>& B) {
  return A.x * B.y - B.x * A.y;
}

template <typename T>
HOST_DEVICE_INLINE void get_rotated_vertices(
    const RotatedBox<T>& box,
    Point<T> (&pts)[4]) {
  // M_PI / 180. == 0.01745329251
  //double theta = box.a * 0.01745329251;
  //MODIFIED
  double theta = box.a;
  T cosTheta2 = (T)cos(theta) * 0.5f;
  T sinTheta2 = (T)sin(theta) * 0.5f;

  // y: top --> down; x: left --> right
  // pts[0].x = box.x_ctr - sinTheta2 * box.h - cosTheta2 * box.w;
  // pts[0].y = box.y_ctr + cosTheta2 * box.h - sinTheta2 * box.w;
  // pts[1].x = box.x_ctr + sinTheta2 * box.h - cosTheta2 * box.w;
  // pts[1].y = box.y_ctr - cosTheta2 * box.h - sinTheta2 * box.w;

  pts[0].x = box.x_ctr + sinTheta2 * box.h + cosTheta2 * box.w;
  pts[0].y = box.y_ctr + cosTheta2 * box.h - sinTheta2 * box.w;
  pts[1].x = box.x_ctr - sinTheta2 * box.h + cosTheta2 * box.w;
  pts[1].y = box.y_ctr - cosTheta2 * box.h - sinTheta2 * box.w;
  pts[2].x = 2 * box.x_ctr - pts[0].x;
  pts[2].y = 2 * box.y_ctr - pts[0].y;
  pts[3].x = 2 * box.x_ctr - pts[1].x;
  pts[3].y = 2 * box.y_ctr - pts[1].y;
}

template <typename T>
HOST_DEVICE_INLINE int get_intersection_points(
    const Point<T> (&pts1)[4],
    const Point<T> (&pts2)[4],
    Point<T> (&intersections)[24]) {
  // Line vector
  // A line from p1 to p2 is: p1 + (p2-p1)*t, t=[0,1]
  Point<T> vec1[4], vec2[4];
  for (int i = 0; i < 4; i++) {
    vec1[i] = pts1[(i + 1) % 4] - pts1[i];
    vec2[i] = pts2[(i + 1) % 4] - pts2[i];
  }

  // Line test - test all line combos for intersection
  int num = 0; // number of intersections
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 4; j++) {
      // Solve for 2x2 Ax=b
      T det = cross_2d<T>(vec2[j], vec1[i]);

      // This takes care of parallel lines
      if (fabs(det) <= 1e-14) {
        continue;
      }

      auto vec12 = pts2[j] - pts1[i];

      T t1 = cross_2d<T>(vec2[j], vec12) / det;
      T t2 = cross_2d<T>(vec1[i], vec12) / det;

      if (t1 >= 0.0f && t1 <= 1.0f && t2 >= 0.0f && t2 <= 1.0f) {
        intersections[num++] = pts1[i] + vec1[i] * t1;
      }
    }
  }

  // Check for vertices of rect1 inside rect2
  {
    const auto& AB = vec2[0];
    const auto& DA = vec2[3];
    auto ABdotAB = dot_2d<T>(AB, AB);
    auto ADdotAD = dot_2d<T>(DA, DA);
    for (int i = 0; i < 4; i++) {
      // assume ABCD is the rectangle, and P is the point to be judged
      // P is inside ABCD iff. P's projection on AB lies within AB
      // and P's projection on AD lies within AD

      auto AP = pts1[i] - pts2[0];

      auto APdotAB = dot_2d<T>(AP, AB);
      auto APdotAD = -dot_2d<T>(AP, DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts1[i];
      }
    }
  }

  // Reverse the check - check for vertices of rect2 inside rect1
  {
    const auto& AB = vec1[0];
    const auto& DA = vec1[3];
    auto ABdotAB = dot_2d<T>(AB, AB);
    auto ADdotAD = dot_2d<T>(DA, DA);
    for (int i = 0; i < 4; i++) {
      auto AP = pts2[i] - pts1[0];

      auto APdotAB = dot_2d<T>(AP, AB);
      auto APdotAD = -dot_2d<T>(AP, DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts2[i];
      }
    }
  }

  return num;
}

template <typename T>
HOST_DEVICE_INLINE int convex_hull_graham(
    const Point<T> (&p)[24],
    const int& num_in,
    Point<T> (&q)[24],
    bool shift_to_zero = false) {
  assert(num_in >= 2);

  // Step 1:
  // Find point with minimum y
  // if more than 1 points have the same minimum y,
  // pick the one with the minimum x.
  int t = 0;
  for (int i = 1; i < num_in; i++) {
    if (p[i].y < p[t].y || (p[i].y == p[t].y && p[i].x < p[t].x)) {
      t = i;
    }
  }
  auto& start = p[t]; // starting point

  // Step 2:
  // Subtract starting point from every points (for sorting in the next step)
  for (int i = 0; i < num_in; i++) {
    q[i] = p[i] - start;
  }

  // Swap the starting point to position 0
  auto tmp = q[0];
  q[0] = q[t];
  q[t] = tmp;

  // Step 3:
  // Sort point 1 ~ num_in according to their relative cross-product values
  // (essentially sorting according to angles)
  // If the angles are the same, sort according to their distance to origin
  T dist[24];
  for (int i = 0; i < num_in; i++) {
    dist[i] = dot_2d<T>(q[i], q[i]);
  }
'''
IOU_ROTATED_HEADER3 = r'''
// Step 4:
  // Make sure there are at least 2 points (that don't overlap with each other)
  // in the stack
  int k; // index of the non-overlapped second point
  for (k = 1; k < num_in; k++) {
    if (dist[k] > 1e-8) {
      break;
    }
  }
  if (k == num_in) {
    // We reach the end, which means the convex hull is just one point
    q[0] = p[t];
    return 1;
  }
  q[1] = q[k];
  int m = 2; // 2 points in the stack
  // Step 5:
  // Finally we can start the scanning process.
  // When a non-convex relationship between the 3 points is found
  // (either concave shape or duplicated points),
  // we pop the previous point from the stack
  // until the 3-point relationship is convex again, or
  // until the stack only contains two points
  for (int i = k + 1; i < num_in; i++) {
    while (m > 1 && cross_2d<T>(q[i] - q[m - 2], q[m - 1] - q[m - 2]) >= 0) {
      m--;
    }
    q[m++] = q[i];
  }

  // Step 6 (Optional):
  // In general sense we need the original coordinates, so we
  // need to shift the points back (reverting Step 2)
  // But if we're only interested in getting the area/perimeter of the shape
  // We can simply return.
  if (!shift_to_zero) {
    for (int i = 0; i < m; i++) {
      q[i] += start;
    }
  }

  return m;
}

template <typename T>
HOST_DEVICE_INLINE T polygon_area(const Point<T> (&q)[24], const int& m) {
  if (m <= 2) {
    return 0;
  }

  T area = 0;
  for (int i = 1; i < m - 1; i++) {
    area += fabs(cross_2d<T>(q[i] - q[0], q[i + 1] - q[0]));
  }

  return area / 2.0;
}

template <typename T>
HOST_DEVICE_INLINE T rotated_boxes_intersection(
    const RotatedBox<T>& box1,
    const RotatedBox<T>& box2) {
  // There are up to 4 x 4 + 4 + 4 = 24 intersections (including dups) returned
  // from rotated_rect_intersection_pts
  Point<T> intersectPts[24], orderedPts[24];

  Point<T> pts1[4];
  Point<T> pts2[4];
  get_rotated_vertices<T>(box1, pts1);
  get_rotated_vertices<T>(box2, pts2);

  int num = get_intersection_points<T>(pts1, pts2, intersectPts);

  if (num <= 2) {
    return 0.0;
  }

  // Convex Hull to order the intersection points in clockwise order and find
  // the contour area.
  int num_convex = convex_hull_graham<T>(intersectPts, num, orderedPts, true);
  return polygon_area<T>(orderedPts, num_convex);
}

} // namespace

template <typename T>
HOST_DEVICE_INLINE T
single_box_iou_rotated(T const* const box1_raw, T const* const box2_raw) {
  // shift center to the middle point to achieve higher precision in result
  RotatedBox<T> box1, box2;


  auto center_shift_x = (box1_raw[0] + box2_raw[0]) / 2.0;
  auto center_shift_y = (box1_raw[1] + box2_raw[1]) / 2.0;
  box1.x_ctr = box1_raw[0] - center_shift_x;
  box1.y_ctr = box1_raw[1] - center_shift_y;
  box1.w = box1_raw[2];
  box1.h = box1_raw[3];
  box1.a = box1_raw[4];
  box2.x_ctr = box2_raw[0] - center_shift_x;
  box2.y_ctr = box2_raw[1] - center_shift_y;
  box2.w = box2_raw[2];
  box2.h = box2_raw[3];
  box2.a = box2_raw[4];

  const T area1 = box1.w * box1.h;
  const T area2 = box2.w * box2.h;
  if (area1 < 1e-14 || area2 < 1e-14) {
    return 0.f;
  }

  const T intersection = rotated_boxes_intersection<T>(box1, box2);
  const T iou = intersection / (area1 + area2 - intersection);
  return iou;
}
'''
IOU_ROTATED_CPU_HEADER = IOU_ROTATED_HEADER1 + r'''
#define HOST_DEVICE
#define HOST_DEVICE_INLINE HOST_DEVICE inline
''' + IOU_ROTATED_HEADER2 + r'''
  // CPU version
  std::sort(
      q + 1, q + num_in, [](const Point<T>& A, const Point<T>& B) -> bool {
        T temp = cross_2d<T>(A, B);
        if (fabs(temp) < 1e-6) {
          return dot_2d<T>(A, A) < dot_2d<T>(B, B);
        } else {
          return temp > 0;
        }
      });
''' + IOU_ROTATED_HEADER3

IOU_ROTATED_CUDA_HEADER = IOU_ROTATED_HEADER1 + r'''
#define HOST_DEVICE __host__ __device__
#define HOST_DEVICE_INLINE HOST_DEVICE __forceinline__
// 2D block with 32 * 16 = 512 threads per block
const int BLOCK_DIM_X = 32;
const int BLOCK_DIM_Y = 16;
''' + IOU_ROTATED_HEADER2 + r'''
  // CUDA version
  // In the future, we can potentially use thrust
  // for sorting here to improve speed (though not guaranteed)
  for (int i = 1; i < num_in - 1; i++) {
    for (int j = i + 1; j < num_in; j++) {
      T crossProduct = cross_2d<T>(q[i], q[j]);
      if ((crossProduct < -1e-6) ||
          (fabs(crossProduct) < 1e-6 && dist[i] > dist[j])) {
        auto q_tmp = q[i];
        q[i] = q[j];
        q[j] = q_tmp;
        auto dist_tmp = dist[i];
        dist[i] = dist[j];
        dist[j] = dist_tmp;
      }
    }
  }
''' + IOU_ROTATED_HEADER3 + r'''
template <typename T>
__global__ void nms_rotated_cuda_kernel(
    const int n_boxes,
    const float iou_threshold,
    const T* dev_boxes,
    unsigned long long* dev_mask) {
  // nms_rotated_cuda_kernel is modified from torchvision's nms_cuda_kernel

  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

  // if (row_start > col_start) return;

  const int row_size =
      min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
      min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  // Compared to nms_cuda_kernel, where each box is represented with 4 values
  // (x1, y1, x2, y2), each rotated box is represented with 5 values
  // (x_center, y_center, width, height, angle_degrees) here.
  __shared__ T block_boxes[threadsPerBlock * 6];
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x * 6 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 0];
    block_boxes[threadIdx.x * 6 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 1];
    block_boxes[threadIdx.x * 6 + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 2];
    block_boxes[threadIdx.x * 6 + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 3];
    block_boxes[threadIdx.x * 6 + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 4];
    block_boxes[threadIdx.x * 6 + 5] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 6 + 5];
  }
  __syncthreads();

  if (threadIdx.x < row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const T* cur_box = dev_boxes + cur_box_idx * 6;
    int i = 0;
    unsigned long long t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (i = start; i < col_size; i++) {
      // Instead of devIoU used by original horizontal nms, here
      // we use the single_box_iou_rotated function from box_iou_rotated_utils.h
      if (single_box_iou_rotated<T>(cur_box, block_boxes + i * 6) >
          iou_threshold) {
        t |= 1ULL << i;
      }
    }
    const int col_blocks = CeilDIV(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t;
  }
}
template <typename T>
__global__ void box_iou_rotated_cuda_kernel(
    const int n_boxes1,
    const int n_boxes2,
    const T* dev_boxes1,
    const T* dev_boxes2,
    T* dev_ious) {
  const int row_start = blockIdx.x * blockDim.x;
  const int col_start = blockIdx.y * blockDim.y;

  const int row_size = min(n_boxes1 - row_start, blockDim.x);
  const int col_size = min(n_boxes2 - col_start, blockDim.y);

  __shared__ float block_boxes1[BLOCK_DIM_X * 5];
  __shared__ float block_boxes2[BLOCK_DIM_Y * 5];

  // It's safe to copy using threadIdx.x since BLOCK_DIM_X >= BLOCK_DIM_Y
  if (threadIdx.x < row_size && threadIdx.y == 0) {
    block_boxes1[threadIdx.x * 5 + 0] =
        dev_boxes1[(row_start + threadIdx.x) * 5 + 0];
    block_boxes1[threadIdx.x * 5 + 1] =
        dev_boxes1[(row_start + threadIdx.x) * 5 + 1];
    block_boxes1[threadIdx.x * 5 + 2] =
        dev_boxes1[(row_start + threadIdx.x) * 5 + 2];
    block_boxes1[threadIdx.x * 5 + 3] =
        dev_boxes1[(row_start + threadIdx.x) * 5 + 3];
    block_boxes1[threadIdx.x * 5 + 4] =
        dev_boxes1[(row_start + threadIdx.x) * 5 + 4];
  }

  if (threadIdx.x < col_size && threadIdx.y == 0) {
    block_boxes2[threadIdx.x * 5 + 0] =
        dev_boxes2[(col_start + threadIdx.x) * 5 + 0];
    block_boxes2[threadIdx.x * 5 + 1] =
        dev_boxes2[(col_start + threadIdx.x) * 5 + 1];
    block_boxes2[threadIdx.x * 5 + 2] =
        dev_boxes2[(col_start + threadIdx.x) * 5 + 2];
    block_boxes2[threadIdx.x * 5 + 3] =
        dev_boxes2[(col_start + threadIdx.x) * 5 + 3];
    block_boxes2[threadIdx.x * 5 + 4] =
        dev_boxes2[(col_start + threadIdx.x) * 5 + 4];
  }
  __syncthreads();

  if (threadIdx.x < row_size && threadIdx.y < col_size) {
    int offset = (row_start + threadIdx.x) * n_boxes2 + col_start + threadIdx.y;
    dev_ious[offset] = single_box_iou_rotated<T>(
        block_boxes1 + threadIdx.x * 5, block_boxes2 + threadIdx.y * 5);
  }
}
'''

IOU_CUDA_SRC = r'''
@alias(boxes1,in0)
@alias(boxes2,in1)
@alias(ious,out0)
int num_boxes1 = boxes1_shape0;
int num_boxes2 = boxes2_shape0;

  if (num_boxes1 > 0 && num_boxes2 > 0) {
    const int blocks_x = CeilDIV(num_boxes1, BLOCK_DIM_X);
    const int blocks_y = CeilDIV(num_boxes2, BLOCK_DIM_Y);

    dim3 blocks(blocks_x, blocks_y);
    dim3 threads(BLOCK_DIM_X, BLOCK_DIM_Y);

    box_iou_rotated_cuda_kernel<<<blocks, threads, 0>>>(
        num_boxes1,
        num_boxes2,
        boxes1_p,
        boxes2_p,
        ious_p);
  }
'''

IOU_CPU_SRC = r'''
@alias(boxes1,in0)
@alias(boxes2,in1)
@alias(ious,out0)
int num_boxes1 = boxes1_shape0;
int num_boxes2 = boxes2_shape0;
int nps = boxes1_shape1;

  for (int i = 0; i < num_boxes1; i++) {
    for (int j = 0; j < num_boxes2; j++) {
      ious_p[i * num_boxes2 + j] = single_box_iou_rotated(boxes1_p+i*nps, boxes2_p+j*nps);
    }
  }
'''


def box_iou_rotated_v1(boxes1, boxes2):
    assert boxes1.dtype == boxes2.dtype
    num_boxes1 = boxes1.shape[0]
    num_boxes2 = boxes2.shape[0]

    ious = jt.code((num_boxes1 * num_boxes2, ),
                   boxes1.dtype, [boxes1, boxes2],
                   cpu_header=IOU_ROTATED_CPU_HEADER,
                   cuda_header=IOU_ROTATED_CUDA_HEADER,
                   cpu_src=IOU_CPU_SRC,
                   cuda_src=IOU_CUDA_SRC)
    ious = ious.reshape(num_boxes1, num_boxes2)

    # same bug will happen when bbox size is to small
    too_small1 = boxes1[:, [2, 3]].min(1)[0] < 0.001
    too_small2 = boxes2[:, [2, 3]].min(1)[0] < 0.001
    if too_small1.any_() or too_small2.any_():
        inds1 = jt.nonzero(too_small1)
        inds2 = jt.nonzero(too_small2)
        ious[inds1, :] = 0.
        ious[:, inds2] = 0.

    return ious


def test_box_iou_rotated():
    boxes1 = jt.array([[0, 0, 1, 1, 0], [0.5, 0.5, 1, 2, 0]])
    boxes2 = jt.array([[0, 0, 1, 1, 0], [0.5, 0.5, 1, 2, 0]])
    ious = box_iou_rotated_v1(boxes1, boxes2)
    print(ious)


if __name__ == '__main__':
    test_box_iou_rotated()
    jt.flags.use_cuda = 1
    test_box_iou_rotated()

jittordet/ops/box_iou_rotated_v1.py
JittorCV/jittordet
import jittor as jt
from jittor import Function

CUDA_HEAD = r'''
template <typename T>
__global__ void convex_sort_kernel(
    const int nbs, const int npts, const int index_size, const bool circular,
    const T* x, const T* y, const T* m, const int* start_index,
    const int* order, int* convex_index) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < nbs) {
    const T* sub_x = x + i * npts;
    const T* sub_y = y + i * npts;
    const T* sub_m = m + i * npts;
    const int* sub_order = order + i * npts;
    const int sub_start_index = start_index[i];

    int* sub_convex_index = convex_index + i * index_size;
    sub_convex_index[0] = sub_start_index;
    int c_i = 0;

    for (int _j = 0; _j < npts; _j++) {
      const int j = sub_order[_j];
      if (j == sub_start_index)continue;
      if (sub_m[j] < 0.5)continue;

      const T x0 = sub_x[j];
      const T y0 = sub_y[j];
      T x1 = sub_x[sub_convex_index[c_i]];
      T y1 = sub_y[sub_convex_index[c_i]];
      T d = (x1 - x0) * (x1 - x0) + (y1 - y0) * (y1 - y0);
      if (d < 0.000001)continue;

      if (c_i < 2) {
	sub_convex_index[++c_i] = j;
      }
      else {
	T x2 = sub_x[sub_convex_index[c_i - 1]];
	T y2 = sub_y[sub_convex_index[c_i - 1]];
	while(1) {
	  T t = (x1 - x2) * (y0 - y2) - (y1 - y2) * (x0 - x2);
	  if (t >= 0) {
	    sub_convex_index[++c_i] = j;
	    break;
	  }
	  else {
	    if (c_i <= 1) {
	      sub_convex_index[c_i] = j;
	      break;
	    }
	    else {
	      c_i--;
	      x1 = sub_x[sub_convex_index[c_i]];
	      y1 = sub_y[sub_convex_index[c_i]];
	      x2 = sub_x[sub_convex_index[c_i - 1]];
	      y2 = sub_y[sub_convex_index[c_i - 1]];
	    }
	  }
	}
      }
    }
    if (circular) sub_convex_index[++c_i] = sub_convex_index[0];
  }
}
'''


def convex_sort_cpu(pts, masks, circular):
    INF = 10000000
    nbs = pts.size(0)
    npts = pts.size(1)
    index_size = npts + 1 if circular else npts

    masks_t = masks.cast(pts.dtype)
    x_t = pts[:, :, 0]
    y_t = pts[:, :, 1]

    masked_y = masks_t * y_t + (1 - masks_t) * INF
    start_index_t, _ = masked_y.argmin(1, keepdims=True)
    start_x = x_t.gather(1, start_index_t)
    start_y = y_t.gather(1, start_index_t)

    pts_cos = (x_t - start_x) / jt.sqrt((x_t - start_x) * (x_t - start_x) +
                                        (y_t - start_y) *
                                        (y_t - start_y) + EPS)
    order_t, _ = pts_cos.argsort(1, descending=True)

    convex_index_t = jt.full((nbs, index_size), -1).int()
    if npts == 0:
        return convex_index_t

    SRC = f"""
    int nbs = {nbs};
    int npts = {npts};
    int index_size = {index_size};
    bool circular = {"true" if circular else "false"};
    """ + r"""
    auto x = in0_p;
    auto y = in1_p;
    auto m = in2_p;
    auto start_index = in3_p;
    auto order = in4_p;
    auto convex_index = out0_p;

    for (int i = 0; i < nbs; i++) {
        auto* sub_x = x + i * npts;
        auto* sub_y = y + i * npts;
        auto* sub_m = m + i * npts;
        int* sub_order = order + i * npts;
        int* sub_convex_index = convex_index + i * index_size;
        int sub_start_index = start_index[i];

        sub_convex_index[0] = sub_start_index;
        int c_i = 0;

        for (int _j = 0; _j < npts; _j++) {
            int j = sub_order[_j];
            if (j == sub_start_index) continue;
            if (sub_m[j] < 0.5) continue;

            auto x0 = sub_x[j];
            auto y0 = sub_y[j];
            auto x1 = sub_x[sub_convex_index[c_i]];
            auto y1 = sub_y[sub_convex_index[c_i]];
            auto d = (x1 - x0) * (x1 - x0) + (y1 - y0) * (y1 - y0);
            if (d < 0.000001)continue;

            if (c_i < 2) {
                sub_convex_index[++c_i] = j;
            }
            else {
                auto x2 = sub_x[sub_convex_index[c_i - 1]];
                auto y2 = sub_y[sub_convex_index[c_i - 1]];
                    while(1) {
                    auto t = (x1 - x2) * (y0 - y2) - (y1 - y2) * (x0 - x2);
                    if (t >= 0) {
                        sub_convex_index[++c_i] = j;
                        break;
                    }
                    else {
                    if (c_i <= 1) {
                        sub_convex_index[c_i] = j;
                        break;
                    }else {
                        c_i--;
                        x1 = sub_x[sub_convex_index[c_i]];
                        y1 = sub_y[sub_convex_index[c_i]];
                        x2 = sub_x[sub_convex_index[c_i - 1]];
                        y2 = sub_y[sub_convex_index[c_i - 1]];
                    }
                }
                }
            }
        }
        if (circular)sub_convex_index[++c_i] = sub_convex_index[0];
    }
    """
    convex_index_t = jt.code(
        outputs=[convex_index_t],
        inputs=[x_t, y_t, masks_t, start_index_t, order_t],
        cpu_src=SRC)
    return convex_index_t


def convex_sort_gpu(pts, masks, circular):
    INF = 10000000
    EPS = 0.000001
    nbs = pts.size(0)
    npts = pts.size(1)
    index_size = npts + 1 if circular else npts

    masks_t = masks.cast(pts.dtype)
    x_t = pts[:, :, 0]
    y_t = pts[:, :, 1]

    masked_y = masks_t * y_t + (1 - masks_t) * INF
    start_index_t, _ = masked_y.argmin(1, keepdims=True)
    start_x = x_t.gather(1, start_index_t)
    start_y = y_t.gather(1, start_index_t)

    pts_cos = (x_t - start_x) / jt.sqrt((x_t - start_x) * (x_t - start_x) +
                                        (y_t - start_y) *
                                        (y_t - start_y) + EPS)
    order_t, _ = pts_cos.argsort(1, descending=True)

    convex_index_t = jt.full((nbs, index_size), -1).int()
    if npts == 0:
        return convex_index_t

    SRC = f"""
    const int nbs = out0_shape0;
    const int index_size = out0_shape1;
    const bool circular = {"true" if circular else "false"};
    const int npts = circular? index_size-1:index_size;
    dim3 blocks((nbs+511)/512);
    dim3 threads(512);
    convex_sort_kernel<<<blocks, threads>>>(nbs, npts, index_size, circular,in0_p, in1_p, in2_p,in3_p, in4_p,out0_p);
    """
    jt.code(
        outputs=[convex_index_t],
        inputs=[x_t, y_t, masks_t, start_index_t, order_t],
        cuda_header=CUDA_HEAD,
        cuda_src=SRC)
    return convex_index_t


def convex_sort(pts, masks, circular=True):
    assert pts.size(0) == masks.size(0) and pts.size(1) == masks.size(1)
    if jt.flags.use_cuda:
        return convex_sort_gpu(pts, masks, circular)
    return convex_sort_cpu(pts, masks, circular)


def test_convex_sort():
    # pts = jt.array([])
    # masks = jt.array([])

    convex_index = convex_sort(pts, masks)
    print(convex_index)


if __name__ == '__main__':
    test_convex_sort()
    jt.flags.use_cuda = 1
    test_convex_sort()

jittordet/ops/convex_sort.py
JittorCV/jittordet
import math

import jittor as jt
from jittor import nn
from jittor.misc import _pair

HEADER = r'''
#undef out
#include<executor.h>
#include <stdio.h>
#include <math.h>
#include <float.h>

#define CUDA_KERNEL_LOOP(i, n)                                 \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
       i += blockDim.x * gridDim.x)

const int CUDA_NUM_THREADS = 1024;
const int kMaxGridNum = 65535;

inline int GET_BLOCKS(const int N)
{
  return std::min(kMaxGridNum, (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS);
}

template <typename scalar_t>
__device__ scalar_t deformable_im2col_bilinear(const scalar_t *bottom_data, const int data_width,
                                               const int height, const int width, scalar_t h, scalar_t w)
{

  int h_low = floor(h);
  int w_low = floor(w);
  int h_high = h_low + 1;
  int w_high = w_low + 1;

  scalar_t lh = h - h_low;
  scalar_t lw = w - w_low;
  scalar_t hh = 1 - lh, hw = 1 - lw;

  scalar_t v1 = 0;
  if (h_low >= 0 && w_low >= 0)
    v1 = bottom_data[h_low * data_width + w_low];
  scalar_t v2 = 0;
  if (h_low >= 0 && w_high <= width - 1)
    v2 = bottom_data[h_low * data_width + w_high];
  scalar_t v3 = 0;
  if (h_high <= height - 1 && w_low >= 0)
    v3 = bottom_data[h_high * data_width + w_low];
  scalar_t v4 = 0;
  if (h_high <= height - 1 && w_high <= width - 1)
    v4 = bottom_data[h_high * data_width + w_high];

  scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;

  scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}

template <typename scalar_t>
__device__ scalar_t get_gradient_weight(scalar_t argmax_h, scalar_t argmax_w,
                                        const int h, const int w, const int height, const int width)
{

  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
  {
    //empty
    return 0;
  }

  int argmax_h_low = floor(argmax_h);
  int argmax_w_low = floor(argmax_w);
  int argmax_h_high = argmax_h_low + 1;
  int argmax_w_high = argmax_w_low + 1;

  scalar_t weight = 0;
  if (h == argmax_h_low && w == argmax_w_low)
    weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);
  if (h == argmax_h_low && w == argmax_w_high)
    weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);
  if (h == argmax_h_high && w == argmax_w_low)
    weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);
  if (h == argmax_h_high && w == argmax_w_high)
    weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);
  return weight;
}

template <typename scalar_t>
__device__ scalar_t get_coordinate_weight(scalar_t argmax_h, scalar_t argmax_w,
                                          const int height, const int width, const scalar_t *im_data,
                                          const int data_width, const int bp_dir)
{

  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
  {
    //empty
    return 0;
  }

  int argmax_h_low = floor(argmax_h);
  int argmax_w_low = floor(argmax_w);
  int argmax_h_high = argmax_h_low + 1;
  int argmax_w_high = argmax_w_low + 1;

  scalar_t weight = 0;

  if (bp_dir == 0)
  {
    if (argmax_h_low >= 0 && argmax_w_low >= 0)
      weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];
    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
      weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];
    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
      weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];
    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
      weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];
  }
  else if (bp_dir == 1)
  {
    if (argmax_h_low >= 0 && argmax_w_low >= 0)
      weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];
    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
      weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];
    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
      weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];
    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
      weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];
  }

  return weight;
}

template <typename scalar_t>
__global__ void deformable_im2col_gpu_kernel(const int n, const scalar_t *data_im, const scalar_t *data_offset,
                                             const int height, const int width, const int kernel_h, const int kernel_w,
                                             const int pad_h, const int pad_w, const int stride_h, const int stride_w,
                                             const int dilation_h, const int dilation_w, const int channel_per_deformable_group,
                                             const int batch_size, const int num_channels, const int deformable_group,
                                             const int height_col, const int width_col,
                                             scalar_t *data_col)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    // index index of output matrix
    const int w_col = index % width_col;
    const int h_col = (index / width_col) % height_col;
    const int b_col = (index / width_col / height_col) % batch_size;
    const int c_im = (index / width_col / height_col) / batch_size;
    const int c_col = c_im * kernel_h * kernel_w;

    // compute deformable group index
    const int deformable_group_index = c_im / channel_per_deformable_group;

    const int h_in = h_col * stride_h - pad_h;
    const int w_in = w_col * stride_w - pad_w;
    scalar_t *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
    //const scalar_t* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
    const scalar_t *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
    const scalar_t *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;

    for (int i = 0; i < kernel_h; ++i)
    {
      for (int j = 0; j < kernel_w; ++j)
      {
        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
        const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
        const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
        scalar_t val = static_cast<scalar_t>(0);
        const scalar_t h_im = h_in + i * dilation_h + offset_h;
        const scalar_t w_im = w_in + j * dilation_w + offset_w;
        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
        {
          //const scalar_t map_h = i * dilation_h + offset_h;
          //const scalar_t map_w = j * dilation_w + offset_w;
          //const int cur_height = height - h_in;
          //const int cur_width = width - w_in;
          //val = deformable_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
          val = deformable_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
        }
        *data_col_ptr = val;
        data_col_ptr += batch_size * height_col * width_col;
      }
    }
  }
}
template <typename scalar_t>
__global__ void deformable_col2im_gpu_kernel(
    const int n, const scalar_t *data_col, const scalar_t *data_offset,
    const int channels, const int height, const int width,
    const int kernel_h, const int kernel_w,
    const int pad_h, const int pad_w,
    const int stride_h, const int stride_w,
    const int dilation_h, const int dilation_w,
    const int channel_per_deformable_group,
    const int batch_size, const int deformable_group,
    const int height_col, const int width_col,
    scalar_t *grad_im,const int whole_size)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    const int j = (index / width_col / height_col / batch_size) % kernel_w;
    const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;
    const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;
    // compute the start and end of the output

    const int deformable_group_index = c / channel_per_deformable_group;

    int w_out = index % width_col;
    int h_out = (index / width_col) % height_col;
    int b = (index / width_col / height_col) % batch_size;
    int w_in = w_out * stride_w - pad_w;
    int h_in = h_out * stride_h - pad_h;

    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) *
                                                        2 * kernel_h * kernel_w * height_col * width_col;
    const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;
    const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;
    const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
    const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
    const scalar_t cur_inv_h_data = h_in + i * dilation_h + offset_h;
    const scalar_t cur_inv_w_data = w_in + j * dilation_w + offset_w;

    const scalar_t cur_top_grad = data_col[index];
    const int cur_h = (int)cur_inv_h_data;
    const int cur_w = (int)cur_inv_w_data;
    for (int dy = -2; dy <= 2; dy++)
    {
      for (int dx = -2; dx <= 2; dx++)
      {
        if (cur_h + dy >= 0 && cur_h + dy < height &&
            cur_w + dx >= 0 && cur_w + dx < width &&
            abs(cur_inv_h_data - (cur_h + dy)) < 1 &&
            abs(cur_inv_w_data - (cur_w + dx)) < 1)
        {
          int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;
          scalar_t weight = get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);
          atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);
        }
      }
    }
  }
}

template <typename scalar_t>
__global__ void deformable_col2im_coord_gpu_kernel(const int n, const scalar_t *data_col,
                                                   const scalar_t *data_im, const scalar_t *data_offset,
                                                   const int channels, const int height, const int width,
                                                   const int kernel_h, const int kernel_w,
                                                   const int pad_h, const int pad_w,
                                                   const int stride_h, const int stride_w,
                                                   const int dilation_h, const int dilation_w,
                                                   const int channel_per_deformable_group,
                                                   const int batch_size, const int offset_channels, const int deformable_group,
                                                   const int height_col, const int width_col, scalar_t *grad_offset)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    scalar_t val = 0;
    int w = index % width_col;
    int h = (index / width_col) % height_col;
    int c = (index / width_col / height_col) % offset_channels;
    int b = (index / width_col / height_col) / offset_channels;
    // compute the start and end of the output

    const int deformable_group_index = c / (2 * kernel_h * kernel_w);
    const int col_step = kernel_h * kernel_w;
    int cnt = 0;
    const scalar_t *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group *
                                                  batch_size * width_col * height_col;
    const scalar_t *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) *
                                                channel_per_deformable_group / kernel_h / kernel_w * height * width;
    const scalar_t *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 *
                                                        kernel_h * kernel_w * height_col * width_col;

    const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;

    for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)
    {
      const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;
      const int bp_dir = offset_c % 2;

      int j = (col_pos / width_col / height_col / batch_size) % kernel_w;
      int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;
      int w_out = col_pos % width_col;
      int h_out = (col_pos / width_col) % height_col;
      int w_in = w_out * stride_w - pad_w;
      int h_in = h_out * stride_h - pad_h;
      const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);
      const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);
      const scalar_t offset_h = data_offset_ptr[data_offset_h_ptr];
      const scalar_t offset_w = data_offset_ptr[data_offset_w_ptr];
      scalar_t inv_h = h_in + i * dilation_h + offset_h;
      scalar_t inv_w = w_in + j * dilation_w + offset_w;
      if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)
      {
        inv_h = inv_w = -2;
      }
      const scalar_t weight = get_coordinate_weight(
          inv_h, inv_w,
          height, width, data_im_ptr + cnt * height * width, width, bp_dir);
      val += weight * data_col_ptr[col_pos];
      cnt += 1;
    }

    grad_offset[index] = val;
  }
}

'''


def deformable_im2col(data_im, data_offset, channels, height, width, ksize_h,
                      ksize_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
                      dilation_w, parallel_imgs, deformable_group,
                      columns_shape):

    src = f"""const int channels = {channels};
    const int height = {height};
    const int width = {width};
    const int ksize_h = {ksize_h};
    const int ksize_w = {ksize_w};
    const int pad_h = {pad_h};
    const int pad_w = {pad_w};
    const int stride_h = {stride_h};
    const int stride_w ={stride_w} ;
    const int dilation_h = {dilation_h};
    const int dilation_w = {dilation_w};
    const int parallel_imgs = {parallel_imgs};
    const int deformable_group = {deformable_group};
    int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
    int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
    int num_kernels = channels * height_col * width_col * parallel_imgs;
    int channel_per_deformable_group = channels / deformable_group;
    cudaMemsetAsync(out0_p,0,out0->size);
    deformable_im2col_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
            num_kernels, in0_p, in1_p, height, width, ksize_h, ksize_w,
            pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
            channel_per_deformable_group, parallel_imgs, channels, deformable_group,
            height_col, width_col, out0_p);

    """
    return jt.code(
        columns_shape,
        data_im.dtype,
        inputs=[data_im, data_offset],
        cuda_header=HEADER,
        cuda_src=src)


def deformable_col2im_coord(data_col, data_im, data_offset, channels, height,
                            width, ksize_h, ksize_w, pad_h, pad_w, stride_h,
                            stride_w, dilation_h, dilation_w, parallel_imgs,
                            deformable_group, grad_offset_shape):
    src = f"""const int channels = {channels};
    const int height = {height};
    const int width = {width};
    const int ksize_h = {ksize_h};
    const int ksize_w = {ksize_w};
    const int pad_h = {pad_h};
    const int pad_w = {pad_w};
    const int stride_h = {stride_h};
    const int stride_w ={stride_w} ;
    const int dilation_h = {dilation_h};
    const int dilation_w = {dilation_w};
    const int parallel_imgs = {parallel_imgs};
    const int deformable_group = {deformable_group};
    int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
    int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
    int num_kernels = height_col * width_col * 2 * ksize_h * ksize_w * deformable_group * parallel_imgs;
    int channel_per_deformable_group = channels * ksize_h * ksize_w / deformable_group;

    cudaMemsetAsync(out0_p,0,out0->size);
    deformable_col2im_coord_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
    num_kernels, in0_p, in1_p, in2_p, channels, height, width,
    ksize_h, ksize_w, pad_h, pad_w, stride_h, stride_w,
    dilation_h, dilation_w, channel_per_deformable_group,
    parallel_imgs, 2 * ksize_h * ksize_w * deformable_group, deformable_group,
    height_col, width_col, out0_p);
    """
    return jt.code(
        grad_offset_shape,
        data_offset.dtype, [data_col, data_im, data_offset],
        cuda_header=HEADER,
        cuda_src=src)


def deformable_col2im(data_col, data_offset, channels, height, width, ksize_h,
                      ksize_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
                      dilation_w, parallel_imgs, deformable_group,
                      grad_im_shape):
    src = f"""const int channels = {channels};
    const int height = {height};
    const int width = {width};
    const int ksize_h = {ksize_h};
    const int ksize_w = {ksize_w};
    const int pad_h = {pad_h};
    const int pad_w = {pad_w};
    const int stride_h = {stride_h};
    const int stride_w ={stride_w} ;
    const int dilation_h = {dilation_h};
    const int dilation_w = {dilation_w};
    const int parallel_imgs = {parallel_imgs};
    const int deformable_group = {deformable_group};
    const int whole_size = {sum(grad_im_shape)};
    int height_col = (height + 2 * pad_h - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
    int width_col = (width + 2 * pad_w - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
    int num_kernels = channels * ksize_h * ksize_w * height_col * width_col * parallel_imgs;
    int channel_per_deformable_group = channels / deformable_group;

    cudaMemsetAsync(out0_p,0,out0->size);
    deformable_col2im_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
            num_kernels, in0_p, in1_p, channels, height, width, ksize_h,
            ksize_w, pad_h, pad_w, stride_h, stride_w,
            dilation_h, dilation_w, channel_per_deformable_group,
            parallel_imgs, deformable_group, height_col, width_col, out0_p,whole_size);
    """
    return jt.code(
        grad_im_shape,
        data_col.dtype, [data_col, data_offset],
        cuda_header=HEADER,
        cuda_src=src)


def deform_conv_forward_cuda(input, weight, offset, kW, kH, dW, dH, padW, padH,
                             dilationW, dilationH, group, deformable_group,
                             im2col_step):

    batchSize, nInputPlane, inputHeight, inputWidth = input.shape

    nOutputPlane = weight.size(0)

    outputWidth = (inputWidth + 2 * padW - (dilationW *
                                            (kW - 1) + 1)) // dW + 1
    outputHeight = (inputHeight + 2 * padH - (dilationH *
                                              (kH - 1) + 1)) // dH + 1

    assert (offset.size(0) == batchSize), 'invalid batch size of offset'

    columns_shape = (nInputPlane * kW * kH,
                     im2col_step * outputHeight * outputWidth)

    input = input.view((batchSize // im2col_step, im2col_step, nInputPlane,
                        inputHeight, inputWidth))
    offset = offset.view(
        (batchSize // im2col_step, im2col_step, deformable_group * 2 * kH * kW,
         outputHeight, outputWidth))

    output_buffer = jt.zeros((batchSize // im2col_step, nOutputPlane,
                              im2col_step * outputHeight, outputWidth),
                             input.dtype)

    output_buffer = output_buffer.view(
        (output_buffer.size(0), group, output_buffer.size(1) // group,
         output_buffer.size(2), output_buffer.size(3)))

    for elt in range(batchSize // im2col_step):
        columns = deformable_im2col(input[elt], offset[elt], nInputPlane,
                                    inputHeight, inputWidth, kH, kW, padH,
                                    padW, dH, dW, dilationH, dilationW,
                                    im2col_step, deformable_group,
                                    columns_shape)

        columns = columns.view(
            (group, columns.size(0) // group, columns.size(1)))
        weight = weight.view((group, weight.size(0) // group, weight.size(1),
                              weight.size(2), weight.size(3)))

        for g in range(group):
            output_buffer[elt, g] = (
                output_buffer[elt, g].flatten(1) +
                jt.matmul(weight[g].flatten(1), columns[g])).view_as(
                    output_buffer[elt, g])

    output_buffer = output_buffer.view(
        (output_buffer.size(0), output_buffer.size(1) * output_buffer.size(2),
         output_buffer.size(3), output_buffer.size(4)))
    output_buffer = output_buffer.view(
        (batchSize // im2col_step, nOutputPlane, im2col_step, outputHeight,
         outputWidth))
    output_buffer = output_buffer.transpose(0, 2, 1, 3, 4)
    output = output_buffer.view(
        (batchSize, nOutputPlane, outputHeight, outputWidth))

    return output


def deform_conv_backward_input_cuda(input, offset, gradOutput, weight, kW, kH,
                                    dW, dH, padW, padH, dilationW, dilationH,
                                    group, deformable_group, im2col_step):

    batchSize, nInputPlane, inputHeight, inputWidth = input.shape
    nOutputPlane = weight.size(0)
    outputWidth = (inputWidth + 2 * padW - (dilationW *
                                            (kW - 1) + 1)) // dW + 1
    outputHeight = (inputHeight + 2 * padH - (dilationH *
                                              (kH - 1) + 1)) // dH + 1

    assert offset.size(0) == batchSize, 'invalid batch size of offset'

    gradInput = jt.zeros_like(input)
    gradOffset = jt.zeros_like(offset)
    columns = jt.zeros(
        (nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth),
        input.dtype)

    gradOutput = gradOutput.view((batchSize // im2col_step, im2col_step,
                                  nOutputPlane, outputHeight, outputWidth))
    gradOutput = gradOutput.transpose(0, 2, 1, 3, 4)

    gradInput = gradInput.view((batchSize // im2col_step, im2col_step,
                                nInputPlane, inputHeight, inputWidth))
    input = input.view((batchSize // im2col_step, im2col_step, nInputPlane,
                        inputHeight, inputWidth))
    gradOffset = gradOffset.view(
        (batchSize // im2col_step, im2col_step, deformable_group * 2 * kH * kW,
         outputHeight, outputWidth))
    offset = offset.view(
        (batchSize // im2col_step, im2col_step, deformable_group * 2 * kH * kW,
         outputHeight, outputWidth))

    for elt in range(batchSize // im2col_step):
        columns = columns.view(
            (group, columns.size(0) // group, columns.size(1)))
        weight = weight.view((group, weight.size(0) // group, weight.size(1),
                              weight.size(2), weight.size(3)))
        gradOutput = gradOutput.view(
            (gradOutput.size(0), group, gradOutput.size(1) // group,
             gradOutput.size(2), gradOutput.size(3), gradOutput.size(4)))

        for g in range(group):
            columns[g] = jt.matmul(weight[g].flatten(1).transpose(1, 0),
                                   gradOutput[elt, g].flatten(1))

        columns = columns.view(
            (columns.size(0) * columns.size(1), columns.size(2)))
        gradOutput = gradOutput.view(
            (gradOutput.size(0), gradOutput.size(1) * gradOutput.size(2),
             gradOutput.size(3), gradOutput.size(4), gradOutput.size(5)))

        gradOffset[elt] = deformable_col2im_coord(
            columns, input[elt], offset[elt], nInputPlane, inputHeight,
            inputWidth, kH, kW, padH, padW, dH, dW, dilationH, dilationW,
            im2col_step, deformable_group, gradOffset[elt].shape)

        gradInput[elt] = deformable_col2im(columns, offset[elt], nInputPlane,
                                           inputHeight, inputWidth, kH, kW,
                                           padH, padW, dH, dW, dilationH,
                                           dilationW, im2col_step,
                                           deformable_group,
                                           gradInput[elt].shape)

    gradInput = gradInput.view(batchSize, nInputPlane, inputHeight, inputWidth)
    gradOffset = gradOffset.view(
        (batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth))
    return gradInput, gradOffset


def deform_conv_backward_parameters_cuda(input, offset, gradOutput, gradWeight,
                                         kW, kH, dW, dH, padW, padH, dilationW,
                                         dilationH, group, deformable_group,
                                         scale, im2col_step):

    batchSize, nInputPlane, inputHeight, inputWidth = input.shape
    nOutputPlane = gradWeight.size(0)
    outputWidth = (inputWidth + 2 * padW - (dilationW *
                                            (kW - 1) + 1)) // dW + 1
    outputHeight = (inputHeight + 2 * padH - (dilationH *
                                              (kH - 1) + 1)) // dH + 1

    assert offset.size(0) == batchSize, 'invalid batch size of offset'

    columns_shape = (nInputPlane * kW * kH,
                     im2col_step * outputHeight * outputWidth)

    gradOutput = gradOutput.view(batchSize // im2col_step, im2col_step,
                                 nOutputPlane, outputHeight, outputWidth)
    gradOutput = gradOutput.transpose(0, 2, 1, 3, 4)

    # gradOutputBuffer = jt.zeros_like(gradOutput)
    # gradOutputBuffer = gradOutputBuffer.view(batchSize // im2col_step, nOutputPlane, im2col_step,outputHeight, outputWidth)
    gradOutputBuffer = gradOutput
    gradOutputBuffer = gradOutputBuffer.view(batchSize // im2col_step,
                                             nOutputPlane,
                                             im2col_step * outputHeight,
                                             outputWidth)

    # gradOutput = gradOutput.transpose(0,2,1,3,4)
    # gradOutput = gradOutput.view(batchSize, nOutputPlane, outputHeight, outputWidth)

    input = input.view(batchSize // im2col_step, im2col_step, nInputPlane,
                       inputHeight, inputWidth)
    offset = offset.view(batchSize // im2col_step, im2col_step,
                         deformable_group * 2 * kH * kW, outputHeight,
                         outputWidth)

    for elt in range(batchSize // im2col_step):
        columns = deformable_im2col(input[elt], offset[elt], nInputPlane,
                                    inputHeight, inputWidth, kH, kW, padH,
                                    padW, dH, dW, dilationH, dilationW,
                                    im2col_step, deformable_group,
                                    columns_shape)

        gradOutputBuffer = gradOutputBuffer.view(
            gradOutputBuffer.size(0), group,
            gradOutputBuffer.size(1) // group, gradOutputBuffer.size(2),
            gradOutputBuffer.size(3))
        columns = columns.view(group,
                               columns.size(0) // group, columns.size(1))
        gradWeight = gradWeight.view(group,
                                     gradWeight.size(0) // group,
                                     gradWeight.size(1), gradWeight.size(2),
                                     gradWeight.size(3))

        for g in range(group):
            gradWeight[g] = (
                gradWeight[g].flatten(1) +
                scale * jt.matmul(gradOutputBuffer[elt, g].flatten(1),
                                  columns[g].transpose(1, 0))).view_as(
                                      gradWeight[g])

        gradOutputBuffer = gradOutputBuffer.view(
            gradOutputBuffer.size(0),
            gradOutputBuffer.size(1) * gradOutputBuffer.size(2),
            gradOutputBuffer.size(3), gradOutputBuffer.size(4))
        columns = columns.view(
            columns.size(0) * columns.size(1), columns.size(2))
        gradWeight = gradWeight.view(
            gradWeight.size(0) * gradWeight.size(1), gradWeight.size(2),
            gradWeight.size(3), gradWeight.size(4))
    return gradWeight


class DeformConvFunction(jt.Function):

    def execute(self,
                input,
                offset,
                weight,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                deformable_groups=1,
                im2col_step=64):
        if input is not None and input.ndim != 4:
            raise ValueError(
                'Expected 4D tensor as input, got {}D tensor instead.'.format(
                    input.ndim))
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.groups = groups
        self.deformable_groups = deformable_groups
        self.im2col_step = im2col_step

        self.saved_tensors = (input, offset, weight)

        output_shape = DeformConvFunction._output_size(input, weight,
                                                       self.padding,
                                                       self.dilation,
                                                       self.stride)

        if not jt.flags.use_cuda:
            raise NotImplementedError
        else:
            cur_im2col_step = min(self.im2col_step, input.shape[0])
            assert (input.shape[0] %
                    cur_im2col_step) == 0, 'im2col step must divide batchsize'
            output = deform_conv_forward_cuda(
                input, weight, offset, weight.size(3), weight.size(2),
                self.stride[1], self.stride[0], self.padding[1],
                self.padding[0], self.dilation[1], self.dilation[0],
                self.groups, self.deformable_groups, cur_im2col_step)
            assert output.shape == output_shape
        return output

    def grad(self, grad_output):
        input, offset, weight = self.saved_tensors

        grad_input = grad_offset = grad_weight = None

        if not jt.flags.use_cuda:
            raise NotImplementedError
        else:
            cur_im2col_step = min(self.im2col_step, input.shape[0])
            assert (input.shape[0] %
                    cur_im2col_step) == 0, 'im2col step must divide batchsize'

            grad_input, grad_offset = deform_conv_backward_input_cuda(
                input, offset, grad_output, weight, weight.size(3),
                weight.size(2), self.stride[1], self.stride[0],
                self.padding[1], self.padding[0], self.dilation[1],
                self.dilation[0], self.groups, self.deformable_groups,
                cur_im2col_step)

            # print(grad_output.sum(),grad_input.sum(),grad_offset.sum())

            grad_weight = jt.zeros_like(weight)
            grad_weight = deform_conv_backward_parameters_cuda(
                input, offset, grad_output, grad_weight, weight.size(3),
                weight.size(2), self.stride[1], self.stride[0],
                self.padding[1], self.padding[0], self.dilation[1],
                self.dilation[0], self.groups, self.deformable_groups, 1,
                cur_im2col_step)

        return grad_input, grad_offset, grad_weight

    @staticmethod
    def _output_size(input, weight, padding, dilation, stride):
        channels = weight.size(0)
        output_size = (input.size(0), channels)
        for d in range(input.ndim - 2):
            in_size = input.size(d + 2)
            pad = padding[d]
            kernel = dilation[d] * (weight.size(d + 2) - 1) + 1
            stride_ = stride[d]
            output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1, )
        if not all(map(lambda s: s > 0, output_size)):
            raise ValueError(
                'convolution input is too small (output would be {})'.format(
                    'x'.join(map(str, output_size))))
        return output_size


deform_conv = DeformConvFunction.apply


class DeformConv(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 deformable_groups=1,
                 bias=False):
        super(DeformConv, self).__init__()

        assert not bias
        assert in_channels % groups == 0, \
            'in_channels {} cannot be divisible by groups {}'.format(
                in_channels, groups)
        assert out_channels % groups == 0, \
            'out_channels {} cannot be divisible by groups {}'.format(
                out_channels, groups)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.groups = groups
        self.deformable_groups = deformable_groups

        self.weight = jt.zeros(
            (out_channels, in_channels // self.groups, *self.kernel_size))

        self.reset_parameters()

    def reset_parameters(self):
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        nn.init.uniform_(self.weight, -stdv, stdv)

    def execute(self, x, offset):
        return deform_conv(x, offset, self.weight, self.stride, self.padding,
                           self.dilation, self.groups, self.deformable_groups)
        # To fix an assert error in deform_conv_cuda.cpp:128
        # input image is smaller than kernel
        # input_pad = (
        #     x.size(2) < self.kernel_size[0] or x.size(3) < self.kernel_size[1])
        # if input_pad:
        #     pad_h = max(self.kernel_size[0] - x.size(2), 0)
        #     pad_w = max(self.kernel_size[1] - x.size(3), 0)
        #     x = F.pad(x, (0, pad_w, 0, pad_h), 'constant', 0).contiguous()
        #     offset = F.pad(offset, (0, pad_w, 0, pad_h), 'constant',
        #                    0).contiguous()
        # out = deform_conv(x, offset, self.weight, self.stride, self.padding,
        #                   self.dilation, self.groups, self.deformable_groups)
        # if input_pad:
        #     out = out[:, :, :out.size(2) - pad_h, :out.size(3) -
        #               pad_w].contiguous()
        # return out


def test_dcn():
    import numpy as np
    dcn = DeformConv(
        256,
        256, (3, 3), (1, 1),
        padding=(1, 1),
        dilation=(1, 1),
        groups=1,
        deformable_groups=1)
    dcn.load_state_dict(jt.load('/home/lxl/workspace/s2anet/dcn_init.pth'))
    # x,offset = jt.load("/home/lxl/workspace/JDet/test_dcn1.pkl")
    x, offset, out = jt.load('/home/lxl/workspace/JDet/test_dcn_final.pkl')
    x = jt.array(x)
    offset = jt.array(offset)

    output = dcn(x, offset)
    # print(out.sum())
    # print(np.abs(output.numpy()-out).max())

    loss = output.sum()
    # print(loss)
    print(output.sum())
    grads = jt.grad(loss, [x, offset, dcn.weight])
    print(grads[0].sum())
    print(grads[1].sum())
    print(grads[2].sum())


if __name__ == '__main__':
    jt.flags.use_cuda = 1
    test_dcn()

jittordet/ops/dcn_v1.py
JittorCV/jittordet
import math

import jittor as jt
import numpy as np
from jdet.utils.registry import HEADS
from jittor import nn
from jittor.misc import _pair

__all__ = ['DCN']


def dcn_v2_conv_forward(input, offset, mask, weight, bias, stride, padding,
                        dilation, deformable_groups):
    kernel_size = weight.shape[2:4]
    batch = input.shape[0]
    channels = input.shape[1]
    height = input.shape[2]
    width = input.shape[3]
    channels_out = weight.shape[0]

    kernel_h = kernel_size[0]
    kernel_w = kernel_size[1]
    stride_h = stride[0]
    stride_w = stride[1]
    pad_h = padding[0]
    pad_w = padding[1]
    dilation_h = dilation[0]
    dilation_w = dilation[1]

    height_out = (height + 2 * pad_h - (dilation_h *
                                        (kernel_h - 1) + 1)) // stride_h + 1
    width_out = (width + 2 * pad_w - (dilation_w *
                                      (kernel_w - 1) + 1)) // stride_w + 1

    ones = jt.ones((batch, height_out, width_out), dtype=input.dtype)
    colums = jt.empty(
        (batch, channels * kernel_h * kernel_w, 1 * height_out * width_out),
        dtype=input.dtype)
    inputs = [input, weight, bias, offset, mask, ones, colums]

    output_shape = (batch, channels_out, height_out, width_out)
    output_type = input.dtype
    output = jt.code(
        output_shape,
        output_type,
        inputs,
        cuda_header=r'''
#undef out
#include<cstdio>
#include<cstring>
#include<algorithm>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <executor.h>
using namespace std;
namespace jittor {
extern cublasHandle_t cublas_handle;
} // jittor
#define CUDA_KERNEL_LOOP(i, n)                          \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x;   \
      i < (n);                                          \
      i += blockDim.x * gridDim.x)
const int CUDA_NUM_THREADS = 1024;
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}
__device__ float dmcn_im2col_bilinear(const float *bottom_data, const int data_width,
                                      const int height, const int width, float h, float w)
{
  int h_low = floor(h);
  int w_low = floor(w);
  int h_high = h_low + 1;
  int w_high = w_low + 1;
  float lh = h - h_low;
  float lw = w - w_low;
  float hh = 1 - lh, hw = 1 - lw;
  float v1 = 0;
  if (h_low >= 0 && w_low >= 0)
    v1 = bottom_data[h_low * data_width + w_low];
  float v2 = 0;
  if (h_low >= 0 && w_high <= width - 1)
    v2 = bottom_data[h_low * data_width + w_high];
  float v3 = 0;
  if (h_high <= height - 1 && w_low >= 0)
    v3 = bottom_data[h_high * data_width + w_low];
  float v4 = 0;
  if (h_high <= height - 1 && w_high <= width - 1)
    v4 = bottom_data[h_high * data_width + w_high];
  float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
  float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}
__global__ void modulated_deformable_im2col_gpu_kernel(const int n,
                                                       const float *data_im, const float *data_offset, const float *data_mask,
                                                       const int height, const int width, const int kernel_h, const int kernel_w,
                                                       const int pad_h, const int pad_w,
                                                       const int stride_h, const int stride_w,
                                                       const int dilation_h, const int dilation_w,
                                                       const int channel_per_deformable_group,
                                                       const int batch_size, const int num_channels, const int deformable_group,
                                                       const int height_col, const int width_col,
                                                       float *data_col)
{
  // launch channels * batch_size * height_col * width_col cores
  CUDA_KERNEL_LOOP(index, n)
  {
    // NOTE(CharlesShang): different from Dai Jifeng's MXNet implementation, col_buffer is of shape (c*kw*kh, N, oh, ow)
    // here columns is of shape (N, c*kw*kh, oh * ow), need to adapt axis
    // index index of output matrix
    const int w_col = index % width_col;
    const int h_col = (index / width_col) % height_col;
    // const int b_col = (index / width_col / height_col) % batch_size;
    const int b_col = (index / width_col / height_col / num_channels) % batch_size;
    // const int c_im = (index / width_col / height_col) / batch_size;
    const int c_im = (index / width_col / height_col) % num_channels;
    // const int c_col = c_im * kernel_h * kernel_w;
    const int c_col = c_im * kernel_h * kernel_w;
    // compute deformable group index
    const int deformable_group_index = c_im / channel_per_deformable_group;
    const int h_in = h_col * stride_h - pad_h;
    const int w_in = w_col * stride_w - pad_w;
    //  float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
    float *data_col_ptr = data_col + ((b_col * num_channels * kernel_w * kernel_h + c_col) * height_col + h_col) * width_col + w_col;
    //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
    const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
    const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
    const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
    for (int i = 0; i < kernel_h; ++i)
    {
      for (int j = 0; j < kernel_w; ++j)
      {
        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
        const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
        const float offset_h = data_offset_ptr[data_offset_h_ptr];
        const float offset_w = data_offset_ptr[data_offset_w_ptr];
        const float mask = data_mask_ptr[data_mask_hw_ptr];
        float val = static_cast<float>(0);
        const float h_im = h_in + i * dilation_h + offset_h;
        const float w_im = w_in + j * dilation_w + offset_w;
        //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
        {
          //const float map_h = i * dilation_h + offset_h;
          //const float map_w = j * dilation_w + offset_w;
          //const int cur_height = height - h_in;
          //const int cur_width = width - w_in;
          //val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
          val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
        }
        *data_col_ptr = val * mask;
        // data_col_ptr += batch_size * height_col * width_col;
        data_col_ptr += height_col * width_col;
      }
    }
  }
}
__global__ void createBatchGemmBuffer(const float **input_b, float **output_b,
                                      float **columns_b, const float **ones_b,
                                      const float **weight_b, const float **bias_b,
                                      float *input, float *output,
                                      float *columns, float *ones,
                                      float *weight, float *bias,
                                      const int input_stride, const int output_stride,
                                      const int columns_stride, const int ones_stride,
                                      const int num_batches)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_batches)
    {
        input_b[idx] = input + idx * input_stride;
        output_b[idx] = output + idx * output_stride;
        columns_b[idx] = columns + idx * columns_stride;
        ones_b[idx] = ones + idx * ones_stride;
        // share weights and bias within a Mini-Batch
        weight_b[idx] = weight;
        bias_b[idx] = bias;
    }
}
void modulated_deformable_im2col_cuda(
  const float* data_im, const float* data_offset, const float* data_mask,
  const int batch_size, const int channels, const int height_im, const int width_im,
  const int height_col, const int width_col, const int kernel_h, const int kernel_w,
  const int pad_h, const int pad_w, const int stride_h, const int stride_w,
  const int dilation_h, const int dilation_w,
  const int deformable_group, float* data_col) {
  // num_axes should be smaller than block size
  const int channel_per_deformable_group = channels / deformable_group;
  const int num_kernels = channels * batch_size * height_col * width_col;
  modulated_deformable_im2col_gpu_kernel<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
      num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kernel_w,
      pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,
      batch_size, channels, deformable_group, height_col, width_col, data_col);
}
    ''',
        cuda_src=f'''
    const int kernel_h = {kernel_h};
    const int kernel_w = {kernel_w};
    const int stride_h = {stride_h};
    const int stride_w = {stride_w};
    const int pad_h = {pad_h};
    const int pad_w = {pad_w};
    const int dilation_h = {dilation_h};
    const int dilation_w = {dilation_w};
    const int deformable_group = {deformable_groups};
''' + r'''
     @alias(input,in0)
    @alias(weight,in1)
    @alias(bias,in2)
    @alias(offset,in3)
    @alias(mask,in4)
    @alias(ones,in5)
    @alias(columns,in6)
    @alias(output,out0)
    const int batch = input_shape0;
    const int channels = input_shape1;
    const int height = input_shape2;
    const int width = input_shape3;
    const int channels_out = weight_shape0;
    const int channels_kernel = weight_shape1;
    const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
    const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
    // prepare for batch-wise computing, which is significantly faster than instance-wise computing
    // when batch size is large.
    // launch batch threads
    int matrices_size = batch * sizeof(float *);
    const float ** input_b;
    float ** output_b;
    float ** columns_b;
    const float ** ones_b;
    const float ** weight_b;
    const float ** bias_b;
    size_t input_b_allocation;
    size_t output_b_allocation;
    size_t columns_b_allocation;
    size_t ones_b_allocation;
    size_t weight_b_allocation;
    size_t bias_b_allocation;
    input_b = (const float **)exe.allocator->alloc(matrices_size, input_b_allocation);
    output_b = (float **)exe.allocator->alloc(matrices_size, output_b_allocation);
    columns_b = (float **)exe.allocator->alloc(matrices_size, columns_b_allocation);
    ones_b = (const float **)exe.allocator->alloc(matrices_size, ones_b_allocation);
    weight_b = (const float **)exe.allocator->alloc(matrices_size, weight_b_allocation);
    bias_b = (const float **)exe.allocator->alloc(matrices_size, bias_b_allocation);
    const int block = 128;
    const int grid = (batch + block - 1) / block;
    createBatchGemmBuffer<<<grid, block>>>(
        input_b, output_b,
        columns_b, ones_b,
        weight_b, bias_b,
        input_p,
        output_p,
        columns_p,
        ones_p,
        weight_p,
        bias_p,
        channels * width * height,
        channels_out * width_out * height_out,
        channels * kernel_h * kernel_w * height_out * width_out,
        height_out * width_out,
        batch);
    long m_ = channels_out;
    long n_ = height_out * width_out;
    long k_ = 1;
    cublasHandle_t& handle = cublas_handle;
    float alpha = 1.0f;
    float beta = 0.0f;
    cublasSgemmBatched(handle,
                            CUBLAS_OP_T,
                            CUBLAS_OP_N,
                            n_,
                            m_,
                            k_,
                            &alpha,
                            ones_b, k_,
                            bias_b, k_,
                            &beta,
                            output_b, n_,
                            batch);
    modulated_deformable_im2col_cuda(input_p,
                                     offset_p,
                                     mask_p,
                                     batch, channels, height, width,
                                     height_out, width_out, kernel_h, kernel_w,
                                     pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
                                     deformable_group,
                                     columns_p);
    long m = channels_out;
    long n = height_out * width_out;
    long k = channels * kernel_h * kernel_w;
    float beta2 = 1.0f;
    cublasSgemmBatched(handle,
                            CUBLAS_OP_N,
                            CUBLAS_OP_N,
                            n,
                            m,
                            k,
                            &alpha,
                            (const float **)columns_b, n,
                            weight_b, k,
                            &beta2,
                            output_b, n,
                            batch);
    exe.allocator->free(input_b, matrices_size, input_b_allocation);
    exe.allocator->free(output_b, matrices_size, output_b_allocation);
    exe.allocator->free(columns_b, matrices_size, columns_b_allocation);
    exe.allocator->free(ones_b, matrices_size, ones_b_allocation);
    exe.allocator->free(weight_b, matrices_size, weight_b_allocation);
    exe.allocator->free(bias_b, matrices_size, bias_b_allocation);
''')
    return output


def dcn_v2_conv_backward(input, offset, mask, weight, bias, grad_output,
                         stride, padding, dilation, deformable_groups):
    kernel_size = weight.shape[2:4]
    batch = input.shape[0]
    channels = input.shape[1]
    height = input.shape[2]
    width = input.shape[3]
    channels_out = weight.shape[0]

    kernel_h = kernel_size[0]
    kernel_w = kernel_size[1]
    stride_h = stride[0]
    stride_w = stride[1]
    pad_h = padding[0]
    pad_w = padding[1]
    dilation_h = dilation[0]
    dilation_w = dilation[1]

    height_out = (height + 2 * pad_h - (dilation_h *
                                        (kernel_h - 1) + 1)) // stride_h + 1
    width_out = (width + 2 * pad_w - (dilation_w *
                                      (kernel_w - 1) + 1)) // stride_w + 1

    ones = jt.ones((batch, height_out, width_out), dtype=input.dtype)
    colums = jt.empty(
        (batch, channels * kernel_h * kernel_w, 1 * height_out * width_out),
        dtype=input.dtype)
    inputs = [input, weight, bias, offset, mask, ones, colums, grad_output]

    output_shape = [
        input.shape, weight.shape, bias.shape, offset.shape, mask.shape
    ]
    output_type = [
        input.dtype, weight.dtype, bias.dtype, offset.dtype, mask.dtype
    ]
    input_grad, weight_grad, bias_grad, offset_grad, mask_grad = jt.code(
        output_shape,
        output_type,
        inputs,
        cuda_header=r'''
#include<cstdio>
#include<cstring>
#include<algorithm>
#include <cuda_runtime.h>
#include <cublas_v2.h>
using namespace std;
namespace jittor {
extern cublasHandle_t cublas_handle;
} // jittor
#define CUDA_KERNEL_LOOP(i, n)                          \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x;   \
      i < (n);                                          \
      i += blockDim.x * gridDim.x)
const int CUDA_NUM_THREADS = 1024;
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}
__device__ float dmcn_im2col_bilinear(const float *bottom_data, const int data_width,
                                      const int height, const int width, float h, float w)
{
  int h_low = floor(h);
  int w_low = floor(w);
  int h_high = h_low + 1;
  int w_high = w_low + 1;
  float lh = h - h_low;
  float lw = w - w_low;
  float hh = 1 - lh, hw = 1 - lw;
  float v1 = 0;
  if (h_low >= 0 && w_low >= 0)
    v1 = bottom_data[h_low * data_width + w_low];
  float v2 = 0;
  if (h_low >= 0 && w_high <= width - 1)
    v2 = bottom_data[h_low * data_width + w_high];
  float v3 = 0;
  if (h_high <= height - 1 && w_low >= 0)
    v3 = bottom_data[h_high * data_width + w_low];
  float v4 = 0;
  if (h_high <= height - 1 && w_high <= width - 1)
    v4 = bottom_data[h_high * data_width + w_high];
  float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
  float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}
__device__ float dmcn_get_gradient_weight(float argmax_h, float argmax_w,
                                          const int h, const int w, const int height, const int width)
{
  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
  {
    //empty
    return 0;
  }
  int argmax_h_low = floor(argmax_h);
  int argmax_w_low = floor(argmax_w);
  int argmax_h_high = argmax_h_low + 1;
  int argmax_w_high = argmax_w_low + 1;
  float weight = 0;
  if (h == argmax_h_low && w == argmax_w_low)
    weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);
  if (h == argmax_h_low && w == argmax_w_high)
    weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);
  if (h == argmax_h_high && w == argmax_w_low)
    weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);
  if (h == argmax_h_high && w == argmax_w_high)
    weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);
  return weight;
}
__device__ float dmcn_get_coordinate_weight(float argmax_h, float argmax_w,
                                            const int height, const int width, const float *im_data,
                                            const int data_width, const int bp_dir)
{
  if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)
  {
    //empty
    return 0;
  }
  int argmax_h_low = floor(argmax_h);
  int argmax_w_low = floor(argmax_w);
  int argmax_h_high = argmax_h_low + 1;
  int argmax_w_high = argmax_w_low + 1;
  float weight = 0;
  if (bp_dir == 0)
  {
    if (argmax_h_low >= 0 && argmax_w_low >= 0)
      weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];
    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
      weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];
    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
      weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];
    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
      weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];
  }
  else if (bp_dir == 1)
  {
    if (argmax_h_low >= 0 && argmax_w_low >= 0)
      weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];
    if (argmax_h_low >= 0 && argmax_w_high <= width - 1)
      weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];
    if (argmax_h_high <= height - 1 && argmax_w_low >= 0)
      weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];
    if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)
      weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];
  }
  return weight;
}
__global__ void modulated_deformable_im2col_gpu_kernel(const int n,
                                                       const float *data_im, const float *data_offset, const float *data_mask,
                                                       const int height, const int width, const int kernel_h, const int kernel_w,
                                                       const int pad_h, const int pad_w,
                                                       const int stride_h, const int stride_w,
                                                       const int dilation_h, const int dilation_w,
                                                       const int channel_per_deformable_group,
                                                       const int batch_size, const int num_channels, const int deformable_group,
                                                       const int height_col, const int width_col,
                                                       float *data_col)
{
  // launch channels * batch_size * height_col * width_col cores
  CUDA_KERNEL_LOOP(index, n)
  {
    // NOTE(CharlesShang): different from Dai Jifeng's MXNet implementation, col_buffer is of shape (c*kw*kh, N, oh, ow)
    // here columns is of shape (N, c*kw*kh, oh * ow), need to adapt axis
    // index index of output matrix
    const int w_col = index % width_col;
    const int h_col = (index / width_col) % height_col;
    // const int b_col = (index / width_col / height_col) % batch_size;
    const int b_col = (index / width_col / height_col / num_channels) % batch_size;
    // const int c_im = (index / width_col / height_col) / batch_size;
    const int c_im = (index / width_col / height_col) % num_channels;
    // const int c_col = c_im * kernel_h * kernel_w;
    const int c_col = c_im * kernel_h * kernel_w;
    // compute deformable group index
    const int deformable_group_index = c_im / channel_per_deformable_group;
    const int h_in = h_col * stride_h - pad_h;
    const int w_in = w_col * stride_w - pad_w;
    //  float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;
    float *data_col_ptr = data_col + ((b_col * num_channels * kernel_w * kernel_h + c_col) * height_col + h_col) * width_col + w_col;
    //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;
    const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;
    const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
    const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
    for (int i = 0; i < kernel_h; ++i)
    {
      for (int j = 0; j < kernel_w; ++j)
      {
        const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;
        const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;
        const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;
        const float offset_h = data_offset_ptr[data_offset_h_ptr];
        const float offset_w = data_offset_ptr[data_offset_w_ptr];
        const float mask = data_mask_ptr[data_mask_hw_ptr];
        float val = static_cast<float>(0);
        const float h_im = h_in + i * dilation_h + offset_h;
        const float w_im = w_in + j * dilation_w + offset_w;
        //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {
        if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)
        {
          //const float map_h = i * dilation_h + offset_h;
          //const float map_w = j * dilation_w + offset_w;
          //const int cur_height = height - h_in;
          //const int cur_width = width - w_in;
          //val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w);
          val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im);
        }
        *data_col_ptr = val * mask;
        // data_col_ptr += batch_size * height_col * width_col;
        data_col_ptr += height_col * width_col;
      }
    }
  }
}
__global__ void modulated_deformable_col2im_gpu_kernel(const int n,
                                                       const float *data_col, const float *data_offset, const float *data_mask,
                                                       const int channels, const int height, const int width,
                                                       const int kernel_h, const int kernel_w,
                                                       const int pad_h, const int pad_w,
                                                       const int stride_h, const int stride_w,
                                                       const int dilation_h, const int dilation_w,
                                                       const int channel_per_deformable_group,
                                                       const int batch_size, const int deformable_group,
                                                       const int height_col, const int width_col,
                                                       float *grad_im)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    const int j = (index / width_col / height_col / batch_size) % kernel_w;
    const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;
    const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;
    // compute the start and end of the output
    const int deformable_group_index = c / channel_per_deformable_group;
    int w_out = index % width_col;
    int h_out = (index / width_col) % height_col;
    int b = (index / width_col / height_col) % batch_size;
    int w_in = w_out * stride_w - pad_w;
    int h_in = h_out * stride_h - pad_h;
    const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
    const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
    const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;
    const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;
    const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out;
    const float offset_h = data_offset_ptr[data_offset_h_ptr];
    const float offset_w = data_offset_ptr[data_offset_w_ptr];
    const float mask = data_mask_ptr[data_mask_hw_ptr];
    const float cur_inv_h_data = h_in + i * dilation_h + offset_h;
    const float cur_inv_w_data = w_in + j * dilation_w + offset_w;
    const float cur_top_grad = data_col[index] * mask;
    const int cur_h = (int)cur_inv_h_data;
    const int cur_w = (int)cur_inv_w_data;
    for (int dy = -2; dy <= 2; dy++)
    {
      for (int dx = -2; dx <= 2; dx++)
      {
        if (cur_h + dy >= 0 && cur_h + dy < height &&
            cur_w + dx >= 0 && cur_w + dx < width &&
            abs(cur_inv_h_data - (cur_h + dy)) < 1 &&
            abs(cur_inv_w_data - (cur_w + dx)) < 1)
        {
          int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;
          float weight = dmcn_get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);
          atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);
        }
      }
    }
  }
}
__global__ void modulated_deformable_col2im_coord_gpu_kernel(const int n,
                                                             const float *data_col, const float *data_im,
                                                             const float *data_offset, const float *data_mask,
                                                             const int channels, const int height, const int width,
                                                             const int kernel_h, const int kernel_w,
                                                             const int pad_h, const int pad_w,
                                                             const int stride_h, const int stride_w,
                                                             const int dilation_h, const int dilation_w,
                                                             const int channel_per_deformable_group,
                                                             const int batch_size, const int offset_channels, const int deformable_group,
                                                             const int height_col, const int width_col,
                                                             float *grad_offset, float *grad_mask)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    float val = 0, mval = 0;
    int w = index % width_col;
    int h = (index / width_col) % height_col;
    int c = (index / width_col / height_col) % offset_channels;
    int b = (index / width_col / height_col) / offset_channels;
    // compute the start and end of the output
    const int deformable_group_index = c / (2 * kernel_h * kernel_w);
    const int col_step = kernel_h * kernel_w;
    int cnt = 0;
    const float *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col;
    const float *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width;
    const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;
    const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;
    const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;
    for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)
    {
      const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;
      const int bp_dir = offset_c % 2;
      int j = (col_pos / width_col / height_col / batch_size) % kernel_w;
      int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;
      int w_out = col_pos % width_col;
      int h_out = (col_pos / width_col) % height_col;
      int w_in = w_out * stride_w - pad_w;
      int h_in = h_out * stride_h - pad_h;
      const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);
      const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);
      const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out);
      const float offset_h = data_offset_ptr[data_offset_h_ptr];
      const float offset_w = data_offset_ptr[data_offset_w_ptr];
      const float mask = data_mask_ptr[data_mask_hw_ptr];
      float inv_h = h_in + i * dilation_h + offset_h;
      float inv_w = w_in + j * dilation_w + offset_w;
      if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)
      {
        inv_h = inv_w = -2;
      }
      else
      {
        mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w);
      }
      const float weight = dmcn_get_coordinate_weight(
          inv_h, inv_w,
          height, width, data_im_ptr + cnt * height * width, width, bp_dir);
      val += weight * data_col_ptr[col_pos] * mask;
      cnt += 1;
    }
    // KERNEL_ASSIGN(grad_offset[index], offset_req, val);
    grad_offset[index] = val;
    if (offset_c % 2 == 0)
      // KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval);
      grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval;
  }
}
void modulated_deformable_im2col_cuda(
  const float* data_im, const float* data_offset, const float* data_mask,
  const int batch_size, const int channels, const int height_im, const int width_im,
  const int height_col, const int width_col, const int kernel_h, const int kernel_w,
  const int pad_h, const int pad_w, const int stride_h, const int stride_w,
  const int dilation_h, const int dilation_w,
  const int deformable_group, float* data_col) {
  // num_axes should be smaller than block size
  const int channel_per_deformable_group = channels / deformable_group;
  const int num_kernels = channels * batch_size * height_col * width_col;
  modulated_deformable_im2col_gpu_kernel
      <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
      num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kernel_w,
      pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,
      batch_size, channels, deformable_group, height_col, width_col, data_col);
}
void modulated_deformable_col2im_cuda(
  const float* data_col, const float* data_offset, const float* data_mask,
  const int batch_size, const int channels, const int height_im, const int width_im,
  const int height_col, const int width_col, const int kernel_h, const int kernel_w,
  const int pad_h, const int pad_w, const int stride_h, const int stride_w,
  const int dilation_h, const int dilation_w,
  const int deformable_group, float* grad_im){
  const int channel_per_deformable_group = channels / deformable_group;
  const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col;
  modulated_deformable_col2im_gpu_kernel
      <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
        num_kernels, data_col, data_offset, data_mask, channels, height_im, width_im,
        kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w,
        dilation_h, dilation_w, channel_per_deformable_group,
        batch_size, deformable_group, height_col, width_col, grad_im);
}
void modulated_deformable_col2im_coord_cuda(
  const float* data_col, const float* data_im, const float* data_offset, const float* data_mask,
  const int batch_size, const int channels, const int height_im, const int width_im,
  const int height_col, const int width_col, const int kernel_h, const int kernel_w,
  const int pad_h, const int pad_w, const int stride_h, const int stride_w,
  const int dilation_h, const int dilation_w,
  const int deformable_group,
  float* grad_offset, float* grad_mask) {
  const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group;
  const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group;
  modulated_deformable_col2im_coord_gpu_kernel
      <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>>(
        num_kernels, data_col, data_im, data_offset, data_mask, channels, height_im, width_im,
        kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
        dilation_h, dilation_w, channel_per_deformable_group,
        batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col,
        grad_offset, grad_mask);
}
    ''',
        cuda_src=f'''
    const int kernel_h = {kernel_h};
    const int kernel_w = {kernel_w};
    const int stride_h = {stride_h};
    const int stride_w = {stride_w};
    const int pad_h = {pad_h};
    const int pad_w = {pad_w};
    const int dilation_h = {dilation_h};
    const int dilation_w = {dilation_w};
    const int deformable_group = {deformable_groups};
    ''' + r'''
    @alias(input,in0)
    @alias(weight,in1)
    @alias(bias,in2)
    @alias(offset,in3)
    @alias(mask,in4)
    @alias(ones,in5)
    @alias(columns,in6)
    @alias(grad_output,in7)
    @alias(grad_input,out0)
    @alias(grad_weight,out1)
    @alias(grad_bias,out2)
    @alias(grad_offset,out3)
    @alias(grad_mask,out4)
    const int batch = input_shape0;
    const int channels = input_shape1;
    const int height = input_shape2;
    const int width = input_shape3;
    const int channels_out = weight_shape0;
    const int channels_kernel = weight_shape1;
    const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
    const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
    for (int b = 0; b < batch; b++)
    {
        auto input_n = input_p+input_stride0*b;
        auto offset_n = offset_p+offset_stride0*b;
        auto mask_n = mask_p+mask_stride0*b;
        auto grad_output_n = grad_output_p+grad_output_stride0*b;
        auto grad_input_n = grad_input_p+grad_input_stride0*b;
        auto grad_offset_n = grad_offset_p+grad_offset_stride0*b;
        auto grad_mask_n = grad_mask_p+grad_mask_stride0*b;
        long m = channels * kernel_h * kernel_w;
        long n = height_out * width_out;
        long k = channels_out;
        float alpha0  = 1.0f;
        float beta0 = 0.0f;
        cublasHandle_t& handle = cublas_handle;
        cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, n, m, k, &alpha0,
                         grad_output_n, n,
                         weight_p, m, &beta0,
                         columns_p, n);
        // gradient w.r.t. input coordinate data
        modulated_deformable_col2im_coord_cuda(columns_p,
                                               input_n,
                                               offset_n,
                                               mask_n,
                                               1, channels, height, width,
                                               height_out, width_out, kernel_h, kernel_w,
                                               pad_h, pad_w, stride_h, stride_w,
                                               dilation_h, dilation_w, deformable_group,
                                               grad_offset_n,
                                               grad_mask_n);
        // gradient w.r.t. input data
        modulated_deformable_col2im_cuda(columns_p,
                                         offset_n,
                                         mask_n,
                                         1, channels, height, width,
                                         height_out, width_out, kernel_h, kernel_w,
                                         pad_h, pad_w, stride_h, stride_w,
                                         dilation_h, dilation_w, deformable_group,
                                         grad_input_n);
        // gradient w.r.t. weight, dWeight should accumulate across the batch and group
        modulated_deformable_im2col_cuda(
                                         input_n,
                                         offset_n,
                                         mask_n,
                                         1, channels, height, width,
                                         height_out, width_out, kernel_h, kernel_w,
                                         pad_h, pad_w, stride_h, stride_w,
                                         dilation_h, dilation_w, deformable_group,
                                         columns_p);
        long m_ = channels_out;
        long n_ = channels * kernel_h * kernel_w;
        long k_ = height_out * width_out;
        float alpha  = 1.0f;
        float beta = 1.0f;
        cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, n_, m_, k_, &alpha,
                         columns_p, k_,
                         grad_output_n, k_, &beta,
                         grad_weight_p, n_);
        //cublasDestroy(handle);
        // gradient w.r.t. bias
        // long m_ = channels_out;
        // long k__ = height_out * width_out;
        cublasSgemv(handle,
                         CUBLAS_OP_T,
                         k_, m_, &alpha,
                         grad_output_n, k_,
                         ones_p, 1, &beta,
                         grad_bias_p, 1);
    }
    ''')
    return input_grad, offset_grad, mask_grad, weight_grad, bias_grad


class DCN_V2_CONV(jt.Function):

    def execute(self, input, offset, mask, weight, bias, stride, padding,
                dilation, deformable_groups):
        self.input = input
        self.offset = offset
        self.mask = mask
        self.weight = weight
        self.bias = bias
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.deformable_groups = deformable_groups
        output = dcn_v2_conv_forward(input, offset, mask, weight, bias, stride,
                                     padding, dilation, deformable_groups)
        return output

    def grad(self, grad_output):
        input_grad, offset_grad, mask_grad, weight_grad, bias_grad = dcn_v2_conv_backward(
            self.input, self.offset, self.mask, self.weight, self.bias,
            grad_output, self.stride, self.padding, self.dilation,
            self.deformable_groups)
        return input_grad, offset_grad, mask_grad, weight_grad, bias_grad, None, None, None, None


dcn_v2_conv = DCN_V2_CONV.apply


def dcn_v2_pooling_forward(input, bbox, trans, spatial_scale, pooled_size,
                           output_dim, no_trans, group_size, part_size,
                           sample_per_part, trans_std):
    channels = input.shape[1]
    num_bbox = bbox.shape[0]
    assert channels == output_dim, 'input channels and output channels must equal'
    pooled_height = pooled_size
    pooled_width = pooled_size
    output_shape = [(num_bbox, output_dim, pooled_height, pooled_width),
                    (num_bbox, output_dim, pooled_height, pooled_width)]
    output_dtypes = [input.dtype, input.dtype]
    inputs = [input, bbox, trans]
    out, top_count = jt.code(
        output_shape,
        output_dtypes,
        inputs,
        cuda_header=r'''
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define CUDA_KERNEL_LOOP(i, n)                        \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
       i < (n);                                       \
       i += blockDim.x * gridDim.x)
const int CUDA_NUM_THREADS = 1024;
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}
__device__ float bilinear_interp(
    const float *data,
    const float x,
    const float y,
    const int width,
    const int height)
{
  int x1 = floor(x);
  int x2 = ceil(x);
  int y1 = floor(y);
  int y2 = ceil(y);
  float dist_x = static_cast<float>(x - x1);
  float dist_y = static_cast<float>(y - y1);
  float value11 = data[y1 * width + x1];
  float value12 = data[y2 * width + x1];
  float value21 = data[y1 * width + x2];
  float value22 = data[y2 * width + x2];
  float value = (1 - dist_x) * (1 - dist_y) * value11 +
            (1 - dist_x) * dist_y * value12 +
            dist_x * (1 - dist_y) * value21 +
            dist_x * dist_y * value22;
  return value;
}
__global__ void DeformablePSROIPoolForwardKernel(
    const int count,
    const float *bottom_data,
    const float spatial_scale,
    const int channels,
    const int height, const int width,
    const int pooled_height, const int pooled_width,
    const float *bottom_rois, const float *bottom_trans,
    const int no_trans,
    const float trans_std,
    const int sample_per_part,
    const int output_dim,
    const int group_size,
    const int part_size,
    const int num_classes,
    const int channels_each_class,
    float *top_data,
    float *top_count)
{
  CUDA_KERNEL_LOOP(index, count)
  {
    // The output is in order (n, ctop, ph, pw)
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int ctop = (index / pooled_width / pooled_height) % output_dim;
    int n = index / pooled_width / pooled_height / output_dim;
    const float *offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    float roi_start_w = static_cast<float>(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
    float roi_start_h = static_cast<float>(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
    float roi_end_w = static_cast<float>(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
    float roi_end_h = static_cast<float>(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
    // Force too small ROIs to be 1x1
    float roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
    float roi_height = max(roi_end_h - roi_start_h, 0.1);
    // Compute w and h at bottom
    float bin_size_h = roi_height / static_cast<float>(pooled_height);
    float bin_size_w = roi_width / static_cast<float>(pooled_width);
    float sub_bin_size_h = bin_size_h / static_cast<float>(sample_per_part);
    float sub_bin_size_w = bin_size_w / static_cast<float>(sample_per_part);
    int part_h = floor(static_cast<float>(ph) / pooled_height * part_size);
    int part_w = floor(static_cast<float>(pw) / pooled_width * part_size);
    int class_id = ctop / channels_each_class;
    float trans_x = no_trans ? static_cast<float>(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;
    float trans_y = no_trans ? static_cast<float>(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;
    float wstart = static_cast<float>(pw) * bin_size_w + roi_start_w;
    wstart += trans_x * roi_width;
    float hstart = static_cast<float>(ph) * bin_size_h + roi_start_h;
    hstart += trans_y * roi_height;
    float sum = 0;
    int count = 0;
    int gw = floor(static_cast<float>(pw) * group_size / pooled_width);
    int gh = floor(static_cast<float>(ph) * group_size / pooled_height);
    gw = min(max(gw, 0), group_size - 1);
    gh = min(max(gh, 0), group_size - 1);
    const float *offset_bottom_data = bottom_data + (roi_batch_ind * channels) * height * width;
    for (int ih = 0; ih < sample_per_part; ih++)
    {
      for (int iw = 0; iw < sample_per_part; iw++)
      {
        float w = wstart + iw * sub_bin_size_w;
        float h = hstart + ih * sub_bin_size_h;
        // bilinear interpolation
        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
        {
          continue;
        }
        w = min(max(w, 0.), width - 1.);
        h = min(max(h, 0.), height - 1.);
        int c = (ctop * group_size + gh) * group_size + gw;
        float val = bilinear_interp(offset_bottom_data + c * height * width, w, h, width, height);
        sum += val;
        count++;
      }
    }
    top_data[index] = count == 0 ? static_cast<float>(0) : sum / count;
    top_count[index] = count;
  }
}
    ''',
        cuda_src=f'''
    const int no_trans = {no_trans};
    const float spatial_scale = {spatial_scale};
    const int output_dim = {output_dim};
    const int group_size = {group_size};
    const int pooled_size = {pooled_size};
    const int part_size = {part_size};
    const int sample_per_part = {sample_per_part};
    const float trans_std = {trans_std};
    ''' + r'''
    @alias(input,in0)
    @alias(bbox,in1)
    @alias(trans,in2)
    @alias(top_count,out1)
  const int batch = input_shape0;
  const int channels = input_shape1;
  const int height = input_shape2;
  const int width = input_shape3;
  const int channels_trans = no_trans ? 2 : trans_shape1;
  const int num_bbox = bbox_shape0;
  auto pooled_height = pooled_size;
  auto pooled_width = pooled_size;
  long out_size = num_bbox * output_dim * pooled_height * pooled_width;
  const int num_classes = no_trans ? 1 : channels_trans / 2;
  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
  long tmp = out_size % 512L==0? out_size/512L :out_size/512L+1L;
  dim3 grid(std::min(tmp, 4096L));
  dim3 block(512);
  DeformablePSROIPoolForwardKernel<<<grid, block>>>(
        out_size,
        input_p,
        spatial_scale,
        channels,
        height, width,
        pooled_height,
        pooled_width,
        bbox_p,
        trans_p,
        no_trans,
        trans_std,
        sample_per_part,
        output_dim,
        group_size,
        part_size,
        num_classes,
        channels_each_class,
        out_p,
        top_count_p);
    ''')
    return out, top_count


def dcn_v2_pooling_backward(grad_output, input, bbox, trans, output_count,
                            no_trans, spatial_scale, output_dim, group_size,
                            pooled_size, part_size, sample_per_part,
                            trans_std):
    output_shape = [input.shape, trans.shape]
    output_dtype = [grad_output.dtype, trans.dtype]
    inputs = [grad_output, input, bbox, trans, output_count]
    input_grad, trans_grad = jt.code(
        output_shape,
        output_dtype,
        inputs,
        cuda_header=r'''
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define CUDA_KERNEL_LOOP(i, n)                        \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
       i < (n);                                       \
       i += blockDim.x * gridDim.x)
const int CUDA_NUM_THREADS = 1024;
inline int GET_BLOCKS(const int N)
{
  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
}
__global__ void DeformablePSROIPoolBackwardAccKernel(
    const int count,
    const float *top_diff,
    const float *top_count,
    const int num_rois,
    const float spatial_scale,
    const int channels,
    const int height, const int width,
    const int pooled_height, const int pooled_width,
    const int output_dim,
    float *bottom_data_diff, float *bottom_trans_diff,
    const float *bottom_data,
    const float *bottom_rois,
    const float *bottom_trans,
    const int no_trans,
    const float trans_std,
    const int sample_per_part,
    const int group_size,
    const int part_size,
    const int num_classes,
    const int channels_each_class)
{
  CUDA_KERNEL_LOOP(index, count)
  {
    // The output is in order (n, ctop, ph, pw)
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int ctop = (index / pooled_width / pooled_height) % output_dim;
    int n = index / pooled_width / pooled_height / output_dim;
    const float *offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    float roi_start_w = static_cast<float>(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
    float roi_start_h = static_cast<float>(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
    float roi_end_w = static_cast<float>(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
    float roi_end_h = static_cast<float>(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
    // Force too small ROIs to be 1x1
    float roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
    float roi_height = max(roi_end_h - roi_start_h, 0.1);
    // Compute w and h at bottom
    float bin_size_h = roi_height / static_cast<float>(pooled_height);
    float bin_size_w = roi_width / static_cast<float>(pooled_width);
    float sub_bin_size_h = bin_size_h / static_cast<float>(sample_per_part);
    float sub_bin_size_w = bin_size_w / static_cast<float>(sample_per_part);
    int part_h = floor(static_cast<float>(ph) / pooled_height * part_size);
    int part_w = floor(static_cast<float>(pw) / pooled_width * part_size);
    int class_id = ctop / channels_each_class;
    float trans_x = no_trans ? static_cast<float>(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;
    float trans_y = no_trans ? static_cast<float>(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;
    float wstart = static_cast<float>(pw) * bin_size_w + roi_start_w;
    wstart += trans_x * roi_width;
    float hstart = static_cast<float>(ph) * bin_size_h + roi_start_h;
    hstart += trans_y * roi_height;
    if (top_count[index] <= 0)
    {
      continue;
    }
    float diff_val = top_diff[index] / top_count[index];
    const float *offset_bottom_data = bottom_data + roi_batch_ind * channels * height * width;
    float *offset_bottom_data_diff = bottom_data_diff + roi_batch_ind * channels * height * width;
    int gw = floor(static_cast<float>(pw) * group_size / pooled_width);
    int gh = floor(static_cast<float>(ph) * group_size / pooled_height);
    gw = min(max(gw, 0), group_size - 1);
    gh = min(max(gh, 0), group_size - 1);
    for (int ih = 0; ih < sample_per_part; ih++)
    {
      for (int iw = 0; iw < sample_per_part; iw++)
      {
        float w = wstart + iw * sub_bin_size_w;
        float h = hstart + ih * sub_bin_size_h;
        // bilinear interpolation
        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
        {
          continue;
        }
        w = min(max(w, 0.), width - 1.);
        h = min(max(h, 0.), height - 1.);
        int c = (ctop * group_size + gh) * group_size + gw;
        // backward on feature
        int x0 = floor(w);
        int x1 = ceil(w);
        int y0 = floor(h);
        int y1 = ceil(h);
        float dist_x = w - x0, dist_y = h - y0;
        float q00 = (1 - dist_x) * (1 - dist_y);
        float q01 = (1 - dist_x) * dist_y;
        float q10 = dist_x * (1 - dist_y);
        float q11 = dist_x * dist_y;
        int bottom_index_base = c * height * width;
        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x0, q00 * diff_val);
        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x0, q01 * diff_val);
        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x1, q10 * diff_val);
        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x1, q11 * diff_val);
        if (no_trans)
        {
          continue;
        }
        float U00 = offset_bottom_data[bottom_index_base + y0 * width + x0];
        float U01 = offset_bottom_data[bottom_index_base + y1 * width + x0];
        float U10 = offset_bottom_data[bottom_index_base + y0 * width + x1];
        float U11 = offset_bottom_data[bottom_index_base + y1 * width + x1];
        float diff_x = (U11 * dist_y + U10 * (1 - dist_y) - U01 * dist_y - U00 * (1 - dist_y)) * trans_std * diff_val;
        diff_x *= roi_width;
        float diff_y = (U11 * dist_x + U01 * (1 - dist_x) - U10 * dist_x - U00 * (1 - dist_x)) * trans_std * diff_val;
        diff_y *= roi_height;
        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w, diff_x);
        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w, diff_y);
      }
    }
  }
}
    ''',
        cuda_src=f'''
    const int no_trans = {no_trans};
    const float spatial_scale = {spatial_scale};
    const int output_dim = {output_dim};
    const int group_size = {group_size};
    const int pooled_size = {pooled_size};
    const int part_size = {part_size};
    const int sample_per_part = {sample_per_part};
    const float trans_std = {trans_std};
    ''' + r'''
    @alias(out_grad,in0)
    @alias(input,in1)
    @alias(bbox,in2)
    @alias(trans,in3)
    @alias(top_count,in4)
    @alias(input_grad,out0)
    @alias(trans_grad,out1)
    const int batch = input_shape0;
  const int channels = input_shape1;
  const int height = input_shape2;
  const int width = input_shape3;
  const int channels_trans = no_trans ? 2 : trans_shape1;
  const int num_bbox = bbox_shape0;
  auto pooled_height = pooled_size;
  auto pooled_width = pooled_size;
  long out_size = num_bbox * output_dim * pooled_height * pooled_width;
  const int num_classes = no_trans ? 1 : channels_trans / 2;
  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
  long tmp = out_size % 512L==0? out_size/512L :out_size/512L+1L;
  dim3 grid(std::min(tmp, 4096L));
  dim3 block(512);
  DeformablePSROIPoolBackwardAccKernel<<<grid, block>>>(
        out_size,
        out_grad_p,
        top_count_p,
        num_bbox,
        spatial_scale,
        channels,
        height,
        width,
        pooled_height,
        pooled_width,
        output_dim,
        input_grad_p,
        trans_grad_p,
        input_p,
        bbox_p,
        trans_p,
        no_trans,
        trans_std,
        sample_per_part,
        group_size,
        part_size,
        num_classes,
        channels_each_class);
    ''')
    return input_grad, trans_grad


class DCN_V2_POOLING(jt.Function):

    def execute(self,
                input,
                rois,
                offset,
                spatial_scale,
                pooled_size,
                output_dim,
                no_trans,
                group_size=1,
                part_size=None,
                sample_per_part=4,
                trans_std=.0):
        self.spatial_scale = spatial_scale
        self.no_trans = int(no_trans)
        self.output_dim = output_dim
        self.group_size = group_size
        self.pooled_size = pooled_size
        self.part_size = pooled_size if part_size is None else part_size
        self.sample_per_part = sample_per_part
        self.trans_std = trans_std
        output, output_count = dcn_v2_pooling_forward(
            input, rois, offset, self.spatial_scale, self.pooled_size,
            self.output_dim, self.no_trans, self.group_size, self.part_size,
            self.sample_per_part, self.trans_std)
        self.output_count = output_count
        self.input = input
        self.rois = rois
        self.offset = offset
        return output

    def grad(self, grad_output):
        grad_input, grad_offset = dcn_v2_pooling_backward(
            grad_output, self.input, self.rois, self.offset, self.output_count,
            self.no_trans, self.spatial_scale, self.output_dim,
            self.group_size, self.pooled_size, self.part_size,
            self.sample_per_part, self.trans_std)
        return grad_input, None, grad_offset, None, None, None, None, None, None, None, None


dcn_v2_pooling = DCN_V2_POOLING.apply


class DeformConv(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 deformable_groups=1,
                 bias=False):
        super(DeformConv, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.deformable_groups = deformable_groups

        self.weight = jt.zeros((out_channels, in_channels, *self.kernel_size))
        if bias:
            self.bias = jt.zeros((out_channels, ))
        else:
            self.bias = np.zeros((out_channels, ))

        self.reset_parameters()

    def reset_parameters(self):
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        nn.init.uniform_(self.weight, -stdv, stdv)

    def execute(self, x, offset):
        assert x.size(2) > self.kernel_size[0] and x.size(
            3) > self.kernel_size[1]
        mask_shape = list(offset.size())
        mask_shape[1] //= 2
        mask = jt.ones(mask_shape, x.dtype)
        return dcn_v2_conv(x, offset, mask, self.weight, jt.array(self.bias),
                           self.stride, self.padding, self.dilation,
                           self.deformable_groups)


class DCNv2(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation=1,
                 deformable_groups=1):
        super(DCNv2, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.deformable_groups = deformable_groups

        self.weight = jt.zeros((out_channels, in_channels, *self.kernel_size))
        self.bias = jt.zeros((out_channels, ))
        self.reset_parameters()

    def reset_parameters(self):
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        nn.init.constant_(self.bias, 0.0)
        nn.init.uniform_(self.weight, -stdv, stdv)

    def execute(self, input, offset, mask):
        assert 2 * self.deformable_groups * self.kernel_size[
            0] * self.kernel_size[1] == offset.shape[1]
        assert self.deformable_groups * self.kernel_size[0] * self.kernel_size[
            1] == mask.shape[1]

        return dcn_v2_conv(input, offset, mask, self.weight, self.bias,
                           self.stride, self.padding, self.dilation,
                           self.deformable_groups)


@HEADS.register_module()
class DCN(DCNv2):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation=1,
                 deformable_groups=1):
        super(DCN, self).__init__(in_channels, out_channels, kernel_size,
                                  stride, padding, dilation, deformable_groups)

        channels_ = self.deformable_groups * 3 * self.kernel_size[
            0] * self.kernel_size[1]
        self.conv_offset_mask = nn.Conv(
            self.in_channels,
            channels_,
            kernel_size=self.kernel_size,
            stride=self.stride,
            padding=self.padding,
            bias=True)
        self.init_offset()

    def init_offset(self):
        nn.init.constant_(self.conv_offset_mask.weight, 0.0)
        nn.init.constant_(self.conv_offset_mask.bias, 0.0)

    def execute(self, input):
        out = self.conv_offset_mask(input)
        o1, o2, mask = jt.chunk(out, 3, dim=1)
        offset = jt.contrib.concat((o1, o2), dim=1)
        mask = jt.sigmoid(mask)
        return dcn_v2_conv(input, offset, mask, self.weight, self.bias,
                           self.stride, self.padding, self.dilation,
                           self.deformable_groups)


class DCNv2Pooling(nn.Module):

    def __init__(self,
                 spatial_scale,
                 pooled_size,
                 output_dim,
                 no_trans,
                 group_size=1,
                 part_size=None,
                 sample_per_part=4,
                 trans_std=.0):
        super(DCNv2Pooling, self).__init__()
        self.spatial_scale = spatial_scale
        self.pooled_size = pooled_size
        self.output_dim = output_dim
        self.no_trans = no_trans
        self.group_size = group_size
        self.part_size = pooled_size if part_size is None else part_size
        self.sample_per_part = sample_per_part
        self.trans_std = trans_std

    def execute(self, input, rois, offset):
        assert input.shape[1] == self.output_dim
        if self.no_trans:
            o_shape = (0, ) + input.shape[1:]
            offset = jt.empty(o_shape)  #jt.zeros_like(input)
        return dcn_v2_pooling(input, rois, offset, self.spatial_scale,
                              self.pooled_size, self.output_dim, self.no_trans,
                              self.group_size, self.part_size,
                              self.sample_per_part, self.trans_std)


class DCNPooling(DCNv2Pooling):

    def __init__(self,
                 spatial_scale,
                 pooled_size,
                 output_dim,
                 no_trans,
                 group_size=1,
                 part_size=None,
                 sample_per_part=4,
                 trans_std=.0,
                 deform_fc_dim=1024):
        super(DCNPooling,
              self).__init__(spatial_scale, pooled_size, output_dim, no_trans,
                             group_size, part_size, sample_per_part, trans_std)

        self.deform_fc_dim = deform_fc_dim

        if not no_trans:
            self.offset_mask_fc = nn.Sequential(
                nn.Linear(
                    self.pooled_size * self.pooled_size * self.output_dim,
                    self.deform_fc_dim), nn.ReLU(),
                nn.Linear(self.deform_fc_dim, self.deform_fc_dim), nn.ReLU(),
                nn.Linear(self.deform_fc_dim,
                          self.pooled_size * self.pooled_size * 3))
            nn.init.constant_(self.offset_mask_fc[4].weight, 0.0)
            nn.init.constant_(self.offset_mask_fc[4].bias, 0.0)

    def execute(self, input, rois):
        o_shape = (0, ) + input.shape[1:]
        offset = jt.empty(o_shape)  #jt.zeros_like(input)

        if not self.no_trans:

            # do roi_align first
            n = rois.shape[0]
            roi = dcn_v2_pooling(
                input,
                rois,
                offset,
                self.spatial_scale,
                self.pooled_size,
                self.output_dim,
                True,  # no trans
                self.group_size,
                self.part_size,
                self.sample_per_part,
                self.trans_std)

            # build mask and offset
            offset_mask = self.offset_mask_fc(roi.view(n, -1))
            offset_mask = offset_mask.view(n, 3, self.pooled_size,
                                           self.pooled_size)
            o1, o2, mask = jt.chunk(offset_mask, 3, dim=1)
            offset = jt.contrib.concat((o1, o2), dim=1)
            mask = jt.sigmoid(mask)

            # do pooling with offset and mask
            return dcn_v2_pooling(
                input, rois, offset, self.spatial_scale, self.pooled_size,
                self.output_dim, self.no_trans, self.group_size,
                self.part_size, self.sample_per_part, self.trans_std) * mask
        # only roi_align
        return dcn_v2_pooling(input, rois, offset, self.spatial_scale,
                              self.pooled_size, self.output_dim, self.no_trans,
                              self.group_size, self.part_size,
                              self.sample_per_part, self.trans_std)


def test_conv():
    import numpy as np
    jt.flags.use_cuda = 1
    input = jt.array(np.random.randn(2, 64, 128, 128).astype(np.float32))
    # wrap all things (offset and mask) in DCN
    dcn = DCN(
        64, 64, kernel_size=(3, 3), stride=1, padding=1, deformable_groups=2)
    # print(dcn.weight.shape, input.shape)
    output = dcn(input)
    print(output)
    print(jt.grad(output.sum(), input))


def test_pool():
    import numpy as np
    jt.flags.use_cuda = 1
    input = jt.array(np.random.randn(2, 32, 64, 64).astype(np.float32))
    batch_inds = jt.array(np.random.randint(2, size=(20, 1)).astype(np.int32))
    x = jt.array(np.random.randint(256, size=(20, 1))).float32()
    y = jt.array(np.random.randint(256, size=(20, 1))).float32()
    w = jt.array(np.random.randint(64, size=(20, 1))).float32()
    h = jt.array(np.random.randint(64, size=(20, 1))).float32()
    rois = jt.contrib.concat((batch_inds, x, y, x + w, y + h), dim=1)
    offset = jt.array(np.random.randn(20, 2, 7, 7).astype(np.float32))

    # normal roi_align
    pooling = DCNv2Pooling(
        spatial_scale=1.0 / 4,
        pooled_size=7,
        output_dim=32,
        no_trans=True,
        group_size=1,
        trans_std=0.1)

    # deformable pooling
    dpooling = DCNv2Pooling(
        spatial_scale=1.0 / 4,
        pooled_size=7,
        output_dim=32,
        no_trans=False,
        group_size=1,
        trans_std=0.1)

    out = pooling(input, rois, offset)
    dout = dpooling(input, rois, offset)
    print(out)
    print(dout)
    print(jt.grad(out.sum(), input))
    print(jt.grad(dout.sum(), input))


if __name__ == '__main__':
    test_conv()
    test_pool()

jittordet/ops/dcn_v2.py
JittorCV/jittordet
import jittor as jt
from jdet.models.utils.weight_init import normal_init
from jittor import nn

HEADER = r"""
#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)

#define THREADS_PER_BLOCK 1024

inline int GET_BLOCKS(const int N) {
    int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    int max_block_num = 65000;
    return std::min(optimal_block_num, max_block_num);
}

template <typename scalar_t>
__device__ scalar_t bilinear_interpolate(const scalar_t* bottom_data,
    const int height, const int width,
    scalar_t y, scalar_t x) {
    // deal with cases that inverse elements are out of feature map boundary
    // if the feature map's size is [height, width], then its valid pixel
    // coordinates range is: x in [0, width-1] y in [0, height-1]
    if (y < -1.0 || y > height || x < -1.0 || x > width) {
        return 0;
    }

    if (y <= 0) y = 0;
    if (x <= 0) x = 0;

    int y_low = (int)y;
    int x_low = (int)x;
    int y_high;
    int x_high;

    if (y_low >= height - 1) {
        y_high = y_low = height - 1;
        y = (scalar_t)y_low;
    }
    else {
        y_high = y_low + 1;
    }

    if (x_low >= width - 1) {
        x_high = x_low = width - 1;
        x = (scalar_t)x_low;
    }
    else {
        x_high = x_low + 1;
    }

    scalar_t ly = y - y_low;
    scalar_t lx = x - x_low;
    scalar_t hy = 1. - ly;
    scalar_t hx = 1. - lx;
    // do bilinear interpolation
    scalar_t lt = bottom_data[y_low * width + x_low];
    scalar_t rt = bottom_data[y_low * width + x_high];
    scalar_t lb = bottom_data[y_high * width + x_low];
    scalar_t rb = bottom_data[y_high * width + x_high];
    scalar_t w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

    scalar_t val = (w1 * lt + w2 * rt + w3 * lb + w4 * rb);

    return val;
}

template <typename scalar_t>
__device__ void bilinear_interpolate_gradient(const int height, const int width,
    scalar_t y, scalar_t x,
    scalar_t& w1, scalar_t& w2,
    scalar_t& w3, scalar_t& w4,
    int& x_low, int& x_high,
    int& y_low, int& y_high) {
    // deal with cases that inverse elements are out of feature map boundary
    if (y < -1.0 || y > height || x < -1.0 || x > width) {
        w1 = w2 = w3 = w4 = 0.;
        x_low = x_high = y_low = y_high = -1;
        return;
    }

    if (y <= 0) y = 0;
    if (x <= 0) x = 0;

    y_low = (int)y;
    x_low = (int)x;

    if (y_low >= height - 1) {
        y_high = y_low = height - 1;
        y = (scalar_t)y_low;
    }
    else {
        y_high = y_low + 1;
    }

    if (x_low >= width - 1) {
        x_high = x_low = width - 1;
        x = (scalar_t)x_low;
    }
    else {
        x_high = x_low + 1;
    }

    scalar_t ly = y - y_low;
    scalar_t lx = x - x_low;
    scalar_t hy = 1. - ly;
    scalar_t hx = 1. - lx;

    w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
    w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
}

template <typename scalar_t>
__global__ void feature_refine_forward_kernel(
    const int nthreads,
    const int points,
    const scalar_t* bottom_data,
    const scalar_t* best_bboxes,  // of shape (n, h, w, 5)
    const float spatial_scale, const int channels, const int height,
    const int width, scalar_t* top_data) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
        // (n, c, h, w) is an element in the aligned output
        int w = index % width;
        int h = (index / width) % height;
        int c = (index / width / height) % channels;
        int n = index / width / height /
            channels;  // refers to the n^th image within a minibatch

        const scalar_t* bbox_offset =
            best_bboxes + ((n * height + h) * width + w) * 5;
        // for rbbox, there are 5 entries: [x_ctr, y_ctr, w, h, ang]
        scalar_t roi_y = bbox_offset[0] * spatial_scale;
        scalar_t roi_x = bbox_offset[1] * spatial_scale;

        scalar_t px[5] = { roi_x, 0, 0, 0, 0 };
        scalar_t py[5] = { roi_y, 0, 0, 0, 0 };

        if (points > 1) {
            scalar_t roi_w = bbox_offset[2] * spatial_scale;
            scalar_t roi_h = bbox_offset[3] * spatial_scale;
            scalar_t roi_a = bbox_offset[4];

            scalar_t w_2 = roi_w / 2, h_2 = roi_h / 2;
            scalar_t cosa = cosf(roi_a), sina = sinf(roi_a);
            scalar_t wx = cosa * w_2, wy = sina * w_2;
            scalar_t hx = -sina * h_2, hy = cosa * h_2;

            px[1] = roi_x + wx + hx; py[1] = roi_y + wy + hy;
            px[2] = roi_x - wx + hx; py[2] = roi_y - wy + hy;
            px[3] = roi_x - wx - hx; py[3] = roi_y - wy - hy;
            px[4] = roi_x + wx - hx; py[4] = roi_y + wy - hy;
        }

        const scalar_t* offset_bottom_data =
            bottom_data + (n * channels + c) * height * width;

        scalar_t output_val = bottom_data[index];
        for (int i = 0; i < points; i++) {
            output_val += bilinear_interpolate<scalar_t>(offset_bottom_data, height,
                width, py[i], px[i]);
        }
        top_data[index] = output_val;
    }
}

template <typename scalar_t>
__global__ void feature_refine_backward_kernel(
    const int nthreads,
    const int points,
    const scalar_t* top_diff,
    const scalar_t* best_bboxes,  // of shape (n, h, w, 5)
    const float spatial_scale, const int channels, const int height,
    const int width, scalar_t* bottom_diff) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
        // (n, c, h, w) is an element in the input diff
        int w = index % width;
        int h = (index / width) % height;
        int c = (index / width / height) % channels;
        int n = index / width / height /
            channels;  // refers to the n^th image within a minibatch

        const scalar_t* bbox_offset =
            best_bboxes + ((n * height + h) * width + w) * 5;
        // for rbbox, there are 5 entries: [x_ctr, y_ctr, w, h, ang]
        scalar_t roi_y = bbox_offset[0] * spatial_scale;
        scalar_t roi_x = bbox_offset[1] * spatial_scale;

        scalar_t px[5] = { roi_x, 0, 0, 0, 0 };
        scalar_t py[5] = { roi_y, 0, 0, 0, 0 };

        if (points > 1) {
            scalar_t roi_w = bbox_offset[2] * spatial_scale;
            scalar_t roi_h = bbox_offset[3] * spatial_scale;
            scalar_t roi_a = bbox_offset[4];

            scalar_t w_2 = roi_w / 2, h_2 = roi_h / 2;
            scalar_t cosa = cosf(roi_a), sina = sinf(roi_a);
            scalar_t wx = cosa * w_2, wy = sina * w_2;
            scalar_t hx = -sina * h_2, hy = cosa * h_2;

            px[1] = roi_x + wx + hx; py[1] = roi_y + wy + hy;
            px[2] = roi_x - wx + hx; py[2] = roi_y - wy + hy;
            px[3] = roi_x - wx - hx; py[3] = roi_y - wy - hy;
            px[4] = roi_x + wx - hx; py[4] = roi_y + wy - hy;
        }

        scalar_t* offset_bottom_diff =
            bottom_diff + (n * channels + c) * height * width;
        scalar_t value_top_diff = top_diff[index];

        atomicAdd(bottom_diff + index, value_top_diff);
        for (int i = 0; i < points; i++) {
            scalar_t w1, w2, w3, w4;
            int x_low, x_high, y_low, y_high;

            bilinear_interpolate_gradient<scalar_t>(height, width, py[i], px[i], w1,
                w2, w3, w4, x_low, x_high, y_low,
                y_high);
            scalar_t g1 = value_top_diff * w1;
            scalar_t g2 = value_top_diff * w2;
            scalar_t g3 = value_top_diff * w3;
            scalar_t g4 = value_top_diff * w4;
            if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
                atomicAdd(offset_bottom_diff + y_low * width + x_low, g1);
                atomicAdd(offset_bottom_diff + y_low * width + x_high, g2);
                atomicAdd(offset_bottom_diff + y_high * width + x_low, g3);
                atomicAdd(offset_bottom_diff + y_high * width + x_high, g4);
            }
        }
    }
}
"""


def feature_refine_forward(features, best_bboxes, spatial_scale, points):
    src = f"""
    const int output_size = {features.numel()};
    feature_refine_forward_kernel<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK >>>(
                    output_size, {points}, in0_p, in1_p, {spatial_scale},in0_shape1, in0_shape2, in0_shape3, out0_p);
    """
    return jt.code(
        features.shape,
        features.dtype, [features, best_bboxes],
        cuda_header=HEADER,
        cuda_src=src)


def feature_refine_backward(top_grad, best_bboxes, spatial_scale, points):
    bottom_grad = jt.zeros_like(top_grad)
    src = f"""
    const int output_size = {top_grad.numel()};
    feature_refine_backward_kernel<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK >>> (
                    output_size, {points}, in0_p, in1_p, {spatial_scale},
                    in0_shape1, in0_shape2, in0_shape3,out0_p);
    """
    return jt.code(
        top_grad.shape,
        top_grad.dtype, [top_grad, best_bboxes],
        cuda_header=HEADER,
        cuda_src=src)


class FeatureRefineFunction(jt.Function):

    def execute(self, features, best_rbboxes, spatial_scale, points=1):
        self.spatial_scale = spatial_scale
        self.points = points
        self.best_rbboxes = best_rbboxes
        assert points in [1, 5]
        output = feature_refine_forward(features, best_rbboxes, spatial_scale,
                                        points)
        return output

    def grad(self, grad_output):
        best_rbboxes = self.best_rbboxes
        points = self.points
        spatial_scale = self.spatial_scale
        grad_input = feature_refine_backward(grad_output, best_rbboxes,
                                             spatial_scale, points)
        return grad_input, None, None, None


feature_refine = FeatureRefineFunction.apply


class FR(nn.Module):

    def __init__(self, spatial_scale, points=1):
        super(FR, self).__init__()
        self.spatial_scale = float(spatial_scale)
        self.points = points

    def execute(self, features, best_rbboxes):
        return feature_refine(features, best_rbboxes, self.spatial_scale,
                              self.points)

    def __repr__(self):
        format_str = self.__class__.__name__
        format_str += '(spatial_scale={}, points={})'.format(
            self.spatial_scale, self.points)
        return format_str


class FeatureRefineModule(nn.Module):

    def __init__(self,
                 in_channels,
                 featmap_strides,
                 conv_cfg=None,
                 norm_cfg=None):
        super(FeatureRefineModule, self).__init__()
        self.in_channels = in_channels
        self.featmap_strides = featmap_strides
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self._init_layers()
        self.init_weights()

    def _init_layers(self):
        self.fr = nn.ModuleList(
            [FR(spatial_scale=1 / s) for s in self.featmap_strides])
        self.conv_5_1 = nn.Conv2d(
            in_channels=self.in_channels,
            out_channels=self.in_channels,
            kernel_size=(5, 1),
            stride=1,
            padding=(2, 0))
        self.conv_1_5 = nn.Conv2d(
            in_channels=self.in_channels,
            out_channels=self.in_channels,
            kernel_size=(1, 5),
            stride=1,
            padding=(0, 2))
        self.conv_1_1 = nn.Conv2d(
            in_channels=self.in_channels,
            out_channels=self.in_channels,
            kernel_size=1)

    def init_weights(self):
        normal_init(self.conv_5_1, std=0.01)
        normal_init(self.conv_1_5, std=0.01)
        normal_init(self.conv_1_1, std=0.01)

    def execute(self, x, best_rbboxes):
        """
        Args:
            x (list[Tensor]):
                feature maps of multiple scales
            best_rbboxes (list[list[Tensor]]):
                best rbboxes of multiple scales of multiple images
        """
        mlvl_rbboxes = [
            jt.concat(best_rbbox) for best_rbbox in zip(*best_rbboxes)
        ]
        out = []
        for x_scale, best_rbboxes_scale, fr_scale in zip(
                x, mlvl_rbboxes, self.fr):
            feat_scale_1 = self.conv_5_1(self.conv_1_5(x_scale))
            feat_scale_2 = self.conv_1_1(x_scale)
            feat_scale = feat_scale_1 + feat_scale_2
            feat_refined_scale = fr_scale(feat_scale, best_rbboxes_scale)
            out.append(x_scale + feat_refined_scale)
        return out


def test():
    import math
    feat_size_x = 32
    feat_size_y = 32
    stride = 8.0
    spatial_scale = 1.0 / stride
    base_size = 4.0 * stride
    num_imgs = 2
    num_chns = 16

    feat = jt.randn((num_imgs, num_chns, feat_size_x, feat_size_y))

    xc, yc = jt.meshgrid(stride * jt.arange(feat_size_x),
                         stride * jt.arange(feat_size_y))
    xc = xc[None, :, :]
    yc = yc[None, :, :]
    xc = xc + base_size * jt.randn(num_imgs, feat_size_x, feat_size_y)
    yc = yc + base_size * jt.randn(num_imgs, feat_size_x, feat_size_y)
    w = base_size * jt.randn(num_imgs, feat_size_x, feat_size_y).exp()
    h = base_size * jt.randn(num_imgs, feat_size_x, feat_size_y).exp()
    a = -math.pi / 2 * jt.rand(num_imgs, feat_size_x, feat_size_y)
    bbbox = jt.stack([xc, yc, w, h, a], dim=-1)

    f1 = FR(spatial_scale, points=1)
    f2 = FR(spatial_scale, points=5)

    o1 = f1(feat, bbbox)
    o2 = f2(feat, bbbox)

    print(o1.sum(), o2.sum())
    print(jt.grad(o1 + o2, feat).sum())


if __name__ == '__main__':
    jt.flags.use_cuda = 1
    test()

jittordet/ops/fr.py
JittorCV/jittordet
import jittor as jt


def nms(boxes, scores, thresh):
    assert boxes.shape[-1] == 4 and len(scores) == len(boxes)
    if scores.ndim == 1:
        scores = scores.unsqueeze(-1)
    dets = jt.concat([boxes, scores], dim=1)
    return jt.nms(dets, thresh)


def multiclass_nms(mlvl_bboxes, mlvl_scores, score_thr, nms, max_per_img):
    """NMS for multi-class bboxes.

    Args:
        multi_bboxes (Var): shape (n, #class*4) or (n, 4)
        multi_scores (Var): shape (n, #class), where the last column
            contains scores of the background class, but this will be ignored.
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold
        max_num (int, optional): if there are more than max_num bboxes after
            NMS, only top max_num will be kept. Default to -1.

    Returns:
        tuple: (dets, labels), Var of shape (k, 5),
            (k), and (k). Dets are boxes with scores. Labels are 0-based.
    """
    boxes = []
    scores = []
    labels = []
    n_class = mlvl_scores.size(1)
    if mlvl_bboxes.shape[1] > 4:
        mlvl_bboxes = mlvl_bboxes.view(mlvl_bboxes.size(0), -1, 4)
    else:
        mlvl_bboxes = mlvl_bboxes.unsqueeze(1)
        mlvl_bboxes = mlvl_bboxes.expand((mlvl_bboxes.size(0), n_class, 4))
    for j in range(0, n_class - 1):
        bbox_j = mlvl_bboxes[:, j, :]
        score_j = mlvl_scores[:, j:j + 1]
        mask = jt.where(score_j > score_thr)[0]
        bbox_j = bbox_j[mask, :]
        score_j = score_j[mask]
        dets = jt.concat([bbox_j, score_j], dim=1)
        keep = jt.nms(dets, nms['thresh'])
        bbox_j = bbox_j[keep]
        score_j = score_j[keep]
        label_j = jt.ones_like(score_j).int32() * j
        boxes.append(bbox_j)
        scores.append(score_j)
        labels.append(label_j)

    boxes = jt.concat(boxes, dim=0)
    scores = jt.concat(scores, dim=0)
    index, _ = jt.argsort(scores, dim=0, descending=True)
    index = index[:max_per_img, 0]
    boxes = jt.concat([boxes, scores], dim=1)[index]
    labels = jt.concat(labels, dim=0).squeeze(1)[index]
    return boxes, labels

jittordet/ops/nms.py
JittorCV/jittordet
import jittor as jt
from shapely.geometry import Polygon

HEADER = r'''
#undef out
#include <executor.h>
#include <vector>
#include <iostream>
#define THCCeilDiv(a,b) ((a + b - 1) / b)
#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
int const threadsPerBlock = sizeof(unsigned long long) * 8;


#define maxn 10
const double eps=1E-8;

__device__ inline int sig(float d){
    return(d>1e-8)-(d<-1e-8);
}

__device__ inline int point_eq(const float2 a, const float2 b) {
    return sig(a.x - b.x) == 0 && sig(a.y - b.y)==0;
}

__device__ inline void point_swap(float2 *a, float2 *b) {
    float2 temp = *a;
    *a = *b;
    *b = temp;
}

__device__ inline void point_reverse(float2 *first, float2* last)
{
    while ((first!=last)&&(first!=--last)) {
        point_swap (first,last);
        ++first;
    }
}

__device__ inline float cross(float2 o,float2 a,float2 b){  //叉积
    return(a.x-o.x)*(b.y-o.y)-(b.x-o.x)*(a.y-o.y);
}
__device__ inline float area(float2* ps,int n){
    ps[n]=ps[0];
    float res=0;
    for(int i=0;i<n;i++){
        res+=ps[i].x*ps[i+1].y-ps[i].y*ps[i+1].x;
    }
    return res/2.0;
}
__device__ inline int lineCross(float2 a,float2 b,float2 c,float2 d,float2&p){
    float s1,s2;
    s1=cross(a,b,c);
    s2=cross(a,b,d);
    if(sig(s1)==0&&sig(s2)==0) return 2;
    if(sig(s2-s1)==0) return 0;
    p.x=(c.x*s2-d.x*s1)/(s2-s1);
    p.y=(c.y*s2-d.y*s1)/(s2-s1);
    return 1;
}

__device__ inline void polygon_cut(float2*p,int&n,float2 a,float2 b, float2* pp){

    int m=0;p[n]=p[0];
    for(int i=0;i<n;i++){
        if(sig(cross(a,b,p[i]))>0) pp[m++]=p[i];
        if(sig(cross(a,b,p[i]))!=sig(cross(a,b,p[i+1])))
            lineCross(a,b,p[i],p[i+1],pp[m++]);
    }
    n=0;
    for(int i=0;i<m;i++)
        if(!i||!(point_eq(pp[i], pp[i-1])))
            p[n++]=pp[i];
    // while(n>1&&p[n-1]==p[0])n--;
    while(n>1&&point_eq(p[n-1], p[0]))n--;
}

//---------------华丽的分隔线-----------------//
//返回三角形oab和三角形ocd的有向交面积,o是原点//
__device__ inline float intersectArea(float2 a,float2 b,float2 c,float2 d){
    float2 o = make_float2(0,0);
    int s1=sig(cross(o,a,b));
    int s2=sig(cross(o,c,d));
    if(s1==0||s2==0)return 0.0;//退化，面积为0
    // if(s1==-1) swap(a,b);
    // if(s2==-1) swap(c,d);
    if (s1 == -1) point_swap(&a, &b);
    if (s2 == -1) point_swap(&c, &d);
    float2 p[10]={o,a,b};
    int n=3;
    float2 pp[maxn];
    polygon_cut(p,n,o,c,pp);
    polygon_cut(p,n,c,d,pp);
    polygon_cut(p,n,d,o,pp);
    float res=fabs(area(p,n));
    if(s1*s2==-1) res=-res;return res;
}
//求两多边形的交面积
__device__ inline float intersectArea(float2*ps1,int n1,float2*ps2,int n2){
    if(area(ps1,n1)<0) point_reverse(ps1,ps1+n1);
    if(area(ps2,n2)<0) point_reverse(ps2,ps2+n2);
    ps1[n1]=ps1[0];
    ps2[n2]=ps2[0];
    float res=0;
    for(int i=0;i<n1;i++){
        for(int j=0;j<n2;j++){
            res+=intersectArea(ps1[i],ps1[i+1],ps2[j],ps2[j+1]);
        }
    }
    return res;//assumeresispositive!
}

// TODO: optimal if by first calculate the iou between two hbbs
__device__ inline float devPolyIoU(float const * const p, float const * const q) {
    float2 ps1[maxn], ps2[maxn];
    int n1 = 4;
    int n2 = 4;
    for (int i = 0; i < 4; i++) {
        ps1[i].x = p[i * 2];
        ps1[i].y = p[i * 2 + 1];

        ps2[i].x = q[i * 2];
        ps2[i].y = q[i * 2 + 1];
    }
    float inter_area = intersectArea(ps1, n1, ps2, n2);
    float union_area = fabs(area(ps1, n1)) + fabs(area(ps2, n2)) - inter_area;
    float iou = 0;
    if (union_area == 0) {
        iou = (inter_area + 1) / (union_area + 1);
    } else {
        iou = inter_area / union_area;
    }
    return iou;
}

__global__ void poly_nms_kernel(const int n_polys, const float nms_overlap_thresh,
                            const float *dev_polys, unsigned long long *dev_mask) {
    const int row_start = blockIdx.y;
    const int col_start = blockIdx.x;

    const int row_size =
            min(n_polys - row_start * threadsPerBlock, threadsPerBlock);
    const int cols_size =
            min(n_polys - col_start * threadsPerBlock, threadsPerBlock);

    __shared__ float block_polys[threadsPerBlock * 9];
    if (threadIdx.x < cols_size) {
        block_polys[threadIdx.x * 9 + 0] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 0];
        block_polys[threadIdx.x * 9 + 1] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 1];
        block_polys[threadIdx.x * 9 + 2] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 2];
        block_polys[threadIdx.x * 9 + 3] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 3];
        block_polys[threadIdx.x * 9 + 4] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 4];
        block_polys[threadIdx.x * 9 + 5] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 5];
        block_polys[threadIdx.x * 9 + 6] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 6];
        block_polys[threadIdx.x * 9 + 7] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 7];
        block_polys[threadIdx.x * 9 + 8] =
            dev_polys[(threadsPerBlock * col_start + threadIdx.x) * 9 + 8];
    }
    __syncthreads();

    if (threadIdx.x < row_size) {
        const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
        const float *cur_box = dev_polys + cur_box_idx * 9;
        int i = 0;
        unsigned long long t = 0;
        int start = 0;
        if (row_start == col_start) {
            start = threadIdx.x + 1;
        }
        for (i = start; i < cols_size; i++) {
            if (devPolyIoU(cur_box, block_polys + i * 9) > nms_overlap_thresh) {
                t |= 1ULL << i;
            }
        }
        const int col_blocks = THCCeilDiv(n_polys, threadsPerBlock);
        dev_mask[cur_box_idx * col_blocks + col_start] = t;
    }
}
'''


def poly_nms(boxes, nms_overlap_thresh):
    assert boxes.ndim == 2 and boxes.shape[1] == 9
    assert jt.flags.use_cuda
    scores = boxes[:, 8]
    order_t, _ = scores.argsort(0, descending=True)
    boxes_sorted = boxes[order_t]
    SRC = f"""
    const float nms_overlap_thresh = {nms_overlap_thresh};
    """ + r"""
    @alias(boxes_sorted,in0)
    @alias(keep,out0)
    cudaMemsetAsync(keep_p,0,keep->size);
    const int boxes_num = boxes_sorted_shape0;
    const int col_blocks = THCCeilDiv(boxes_num, threadsPerBlock);

    int matrices_size = boxes_num * col_blocks * sizeof(unsigned long long);
    size_t mask_allocation;
    unsigned long long* mask_p = (unsigned long long*)exe.allocator->alloc(matrices_size, mask_allocation);
    dim3 blocks(THCCeilDiv(boxes_num, threadsPerBlock),
                THCCeilDiv(boxes_num, threadsPerBlock));
    dim3 threads(threadsPerBlock);
    poly_nms_kernel<<<blocks, threads, 0>>>(boxes_num,
                                        nms_overlap_thresh,
                                        boxes_sorted_p,
                                        mask_p);
    checkCudaErrors(cudaDeviceSynchronize());

    std::vector<unsigned long long> remv(col_blocks);

    memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);
    for (int i = 0; i < boxes_num; i++) {
        int nblock = i / threadsPerBlock;
        int inblock = i % threadsPerBlock;

        if (!(remv[nblock] & (1ULL << inblock))) {
            keep_p[i] = true;
            unsigned long long* p = mask_p + i * col_blocks;
            for (int j = nblock; j < col_blocks; j++) {
                remv[j] |= p[j];
            }
        }
    }
    exe.allocator->free(mask_p, matrices_size, mask_allocation);
    """
    keep = jt.code((boxes.shape[0], ),
                   'bool', [boxes_sorted],
                   cuda_header=HEADER,
                   cuda_src=SRC)
    return order_t[keep]


def multiclass_poly_nms(bboxes, scores, labels, thresh):
    max_coordinate = bboxes.max() - bboxes.min()
    offsets = labels * (max_coordinate + 1)
    bboxes_for_nms = bboxes + offsets[:, None]

    keep = poly_nms(
        jt.concat([bboxes_for_nms, scores[:, None]], dim=1), thresh)

    bboxes = bboxes[keep]
    scores = scores[keep]
    labels = labels[keep]
    dets = jt.concat([bboxes, scores[:, None]], dim=1)
    return dets, labels


def iou_poly(poly1, poly2):
    poly1 = Polygon(poly1.reshape(4, 2))
    poly2 = Polygon(poly2.reshape(4, 2))
    inter_area = poly1.intersection(poly2).area
    iou = inter_area / max(poly1.area + poly2.area - inter_area, 0.01)
    return iou


def nms_poly_numpy(dets, thresh):
    scores = dets[:, 8]
    polys = dets[:, :8]
    areas = []
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        ovr = []
        i = order[0]
        keep.append(i)
        for j in range(order.size - 1):
            iou = iou_poly(polys[i], polys[order[j + 1]])
            ovr.append(iou)
        ovr = np.array(ovr)
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep


def test2(i):
    import pickle

    import numpy as np

    bboxes, scores, labels = pickle.load(
        open(f'/home/lxl/diff/poly_nms_in_{i}.pkl', 'rb'))
    bboxes_j, scores_j, keep_j = multiclass_poly_nms(
        jt.array(bboxes), jt.array(scores), jt.array(labels), 0.1)

    bboxes_t, scores_t, keep_t = pickle.load(
        open(f'/home/lxl/diff/poly_nms_out_{i}.pkl', 'rb'))
    # print(keep_t)
    # print(keep_j)

    # print(keep_j,np.sort(keep_t))
    # np.testing.assert_allclose(bboxes_j.numpy(),bboxes_t,rtol=1e-05)
    # np.testing.assert_allclose(scores_j.numpy(),scores_t,rtol=1e-05)
    # np.testing.assert_allclose(keep_j.numpy(),keep_t,rtol=1e-05)

    if keep_j.numpy().shape != keep_t.shape or not np.allclose(
            keep_j.numpy(), keep_t, rtol=1e-05):
        print(i, keep_j.shape, keep_t.shape)
        print(keep_j.numpy())
        print(keep_t)


if __name__ == '__main__':
    jt.flags.use_cuda = 1
    # test2(19)
    # for i in range(1000):
    #     test2(i)

jittordet/ops/nms_poly.py
JittorCV/jittordet
import jittor as jt

ML_NMS_ROTATED_HEADER1 = r'''
// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
#undef out
#include <cassert>
#include <cmath>
#include <executor.h>
#include <algorithm>
#define CeilDIV(a,b) ((a+b-1)/b)
'''

ML_NMS_ROTATED_HEADER2 = r'''
namespace {
int const threadsPerBlock = sizeof(unsigned long long) * 8;

template <typename T>
struct RotatedBox {
  T x_ctr, y_ctr, w, h, a;
};

template <typename T>
struct Point {
  T x, y;
  HOST_DEVICE_INLINE Point(const T& px = 0, const T& py = 0) : x(px), y(py) {}
  HOST_DEVICE_INLINE Point operator+(const Point& p) const {
    return Point(x + p.x, y + p.y);
  }
  HOST_DEVICE_INLINE Point& operator+=(const Point& p) {
    x += p.x;
    y += p.y;
    return *this;
  }
  HOST_DEVICE_INLINE Point operator-(const Point& p) const {
    return Point(x - p.x, y - p.y);
  }
  HOST_DEVICE_INLINE Point operator*(const T coeff) const {
    return Point(x * coeff, y * coeff);
  }
};

template <typename T>
HOST_DEVICE_INLINE T dot_2d(const Point<T>& A, const Point<T>& B) {
  return A.x * B.x + A.y * B.y;
}

template <typename T>
HOST_DEVICE_INLINE T cross_2d(const Point<T>& A, const Point<T>& B) {
  return A.x * B.y - B.x * A.y;
}

template <typename T>
HOST_DEVICE_INLINE void get_rotated_vertices(
    const RotatedBox<T>& box,
    Point<T> (&pts)[4]) {
  // M_PI / 180. == 0.01745329251
  //double theta = box.a * 0.01745329251;
  //MODIFIED
  double theta = box.a;
  T cosTheta2 = (T)cos(theta) * 0.5f;
  T sinTheta2 = (T)sin(theta) * 0.5f;

  // y: top --> down; x: left --> right
  pts[0].x = box.x_ctr - sinTheta2 * box.h - cosTheta2 * box.w;
  pts[0].y = box.y_ctr + cosTheta2 * box.h - sinTheta2 * box.w;
  pts[1].x = box.x_ctr + sinTheta2 * box.h - cosTheta2 * box.w;
  pts[1].y = box.y_ctr - cosTheta2 * box.h - sinTheta2 * box.w;
  pts[2].x = 2 * box.x_ctr - pts[0].x;
  pts[2].y = 2 * box.y_ctr - pts[0].y;
  pts[3].x = 2 * box.x_ctr - pts[1].x;
  pts[3].y = 2 * box.y_ctr - pts[1].y;
}

template <typename T>
HOST_DEVICE_INLINE int get_intersection_points(
    const Point<T> (&pts1)[4],
    const Point<T> (&pts2)[4],
    Point<T> (&intersections)[24]) {
  // Line vector
  // A line from p1 to p2 is: p1 + (p2-p1)*t, t=[0,1]
  Point<T> vec1[4], vec2[4];
  for (int i = 0; i < 4; i++) {
    vec1[i] = pts1[(i + 1) % 4] - pts1[i];
    vec2[i] = pts2[(i + 1) % 4] - pts2[i];
  }

  // Line test - test all line combos for intersection
  int num = 0; // number of intersections
  for (int i = 0; i < 4; i++) {
    for (int j = 0; j < 4; j++) {
      // Solve for 2x2 Ax=b
      T det = cross_2d<T>(vec2[j], vec1[i]);

      // This takes care of parallel lines
      if (fabs(det) <= 1e-14) {
        continue;
      }

      auto vec12 = pts2[j] - pts1[i];

      T t1 = cross_2d<T>(vec2[j], vec12) / det;
      T t2 = cross_2d<T>(vec1[i], vec12) / det;

      if (t1 >= 0.0f && t1 <= 1.0f && t2 >= 0.0f && t2 <= 1.0f) {
        intersections[num++] = pts1[i] + vec1[i] * t1;
      }
    }
  }

  // Check for vertices of rect1 inside rect2
  {
    const auto& AB = vec2[0];
    const auto& DA = vec2[3];
    auto ABdotAB = dot_2d<T>(AB, AB);
    auto ADdotAD = dot_2d<T>(DA, DA);
    for (int i = 0; i < 4; i++) {
      // assume ABCD is the rectangle, and P is the point to be judged
      // P is inside ABCD iff. P's projection on AB lies within AB
      // and P's projection on AD lies within AD

      auto AP = pts1[i] - pts2[0];

      auto APdotAB = dot_2d<T>(AP, AB);
      auto APdotAD = -dot_2d<T>(AP, DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts1[i];
      }
    }
  }

  // Reverse the check - check for vertices of rect2 inside rect1
  {
    const auto& AB = vec1[0];
    const auto& DA = vec1[3];
    auto ABdotAB = dot_2d<T>(AB, AB);
    auto ADdotAD = dot_2d<T>(DA, DA);
    for (int i = 0; i < 4; i++) {
      auto AP = pts2[i] - pts1[0];

      auto APdotAB = dot_2d<T>(AP, AB);
      auto APdotAD = -dot_2d<T>(AP, DA);

      if ((APdotAB >= 0) && (APdotAD >= 0) && (APdotAB <= ABdotAB) &&
          (APdotAD <= ADdotAD)) {
        intersections[num++] = pts2[i];
      }
    }
  }

  return num;
}

template <typename T>
HOST_DEVICE_INLINE int convex_hull_graham(
    const Point<T> (&p)[24],
    const int& num_in,
    Point<T> (&q)[24],
    bool shift_to_zero = false) {
  assert(num_in >= 2);

  // Step 1:
  // Find point with minimum y
  // if more than 1 points have the same minimum y,
  // pick the one with the minimum x.
  int t = 0;
  for (int i = 1; i < num_in; i++) {
    if (p[i].y < p[t].y || (p[i].y == p[t].y && p[i].x < p[t].x)) {
      t = i;
    }
  }
  auto& start = p[t]; // starting point

  // Step 2:
  // Subtract starting point from every points (for sorting in the next step)
  for (int i = 0; i < num_in; i++) {
    q[i] = p[i] - start;
  }

  // Swap the starting point to position 0
  auto tmp = q[0];
  q[0] = q[t];
  q[t] = tmp;

  // Step 3:
  // Sort point 1 ~ num_in according to their relative cross-product values
  // (essentially sorting according to angles)
  // If the angles are the same, sort according to their distance to origin
  T dist[24];
  for (int i = 0; i < num_in; i++) {
    dist[i] = dot_2d<T>(q[i], q[i]);
  }
'''
ML_NMS_ROTATED_HEADER3 = r'''
// Step 4:
  // Make sure there are at least 2 points (that don't overlap with each other)
  // in the stack
  int k; // index of the non-overlapped second point
  for (k = 1; k < num_in; k++) {
    if (dist[k] > 1e-8) {
      break;
    }
  }
  if (k == num_in) {
    // We reach the end, which means the convex hull is just one point
    q[0] = p[t];
    return 1;
  }
  q[1] = q[k];
  int m = 2; // 2 points in the stack
  // Step 5:
  // Finally we can start the scanning process.
  // When a non-convex relationship between the 3 points is found
  // (either concave shape or duplicated points),
  // we pop the previous point from the stack
  // until the 3-point relationship is convex again, or
  // until the stack only contains two points
  for (int i = k + 1; i < num_in; i++) {
    while (m > 1 && cross_2d<T>(q[i] - q[m - 2], q[m - 1] - q[m - 2]) >= 0) {
      m--;
    }
    q[m++] = q[i];
  }

  // Step 6 (Optional):
  // In general sense we need the original coordinates, so we
  // need to shift the points back (reverting Step 2)
  // But if we're only interested in getting the area/perimeter of the shape
  // We can simply return.
  if (!shift_to_zero) {
    for (int i = 0; i < m; i++) {
      q[i] += start;
    }
  }

  return m;
}

template <typename T>
HOST_DEVICE_INLINE T polygon_area(const Point<T> (&q)[24], const int& m) {
  if (m <= 2) {
    return 0;
  }

  T area = 0;
  for (int i = 1; i < m - 1; i++) {
    area += fabs(cross_2d<T>(q[i] - q[0], q[i + 1] - q[0]));
  }

  return area / 2.0;
}

template <typename T>
HOST_DEVICE_INLINE T rotated_boxes_intersection(
    const RotatedBox<T>& box1,
    const RotatedBox<T>& box2) {
  // There are up to 4 x 4 + 4 + 4 = 24 intersections (including dups) returned
  // from rotated_rect_intersection_pts
  Point<T> intersectPts[24], orderedPts[24];

  Point<T> pts1[4];
  Point<T> pts2[4];
  get_rotated_vertices<T>(box1, pts1);
  get_rotated_vertices<T>(box2, pts2);

  int num = get_intersection_points<T>(pts1, pts2, intersectPts);

  if (num <= 2) {
    return 0.0;
  }

  // Convex Hull to order the intersection points in clockwise order and find
  // the contour area.
  int num_convex = convex_hull_graham<T>(intersectPts, num, orderedPts, true);
  return polygon_area<T>(orderedPts, num_convex);
}

} // namespace

template <typename T>
HOST_DEVICE_INLINE T
single_box_iou_rotated(T const* const box1_raw, T const* const box2_raw) {
  // we dont calculate IoU if two bboxes belong to two classes and set it to zero
  // box: [x,y,w,h,a,l]
  if (BOX_LENGTH==6 && box1_raw[5] != box2_raw[5])
    return 0.0;
  // shift center to the middle point to achieve higher precision in result
  RotatedBox<T> box1, box2;

  auto center_shift_x = (box1_raw[0] + box2_raw[0]) / 2.0;
  auto center_shift_y = (box1_raw[1] + box2_raw[1]) / 2.0;
  box1.x_ctr = box1_raw[0] - center_shift_x;
  box1.y_ctr = box1_raw[1] - center_shift_y;
  box1.w = box1_raw[2];
  box1.h = box1_raw[3];
  box1.a = box1_raw[4];
  box2.x_ctr = box2_raw[0] - center_shift_x;
  box2.y_ctr = box2_raw[1] - center_shift_y;
  box2.w = box2_raw[2];
  box2.h = box2_raw[3];
  box2.a = box2_raw[4];

  const T area1 = box1.w * box1.h;
  const T area2 = box2.w * box2.h;
  if (area1 < 1e-14 || area2 < 1e-14) {
    return 0.f;
  }

  const T intersection = rotated_boxes_intersection<T>(box1, box2);
  const T iou = intersection / (area1 + area2 - intersection);
  return iou;
}
'''
ML_NMS_ROTATED_CPU_HEADER = ML_NMS_ROTATED_HEADER1 + r'''
#define HOST_DEVICE
#define HOST_DEVICE_INLINE HOST_DEVICE inline
''' + ML_NMS_ROTATED_HEADER2 + r'''
  // CPU version
  std::sort(
      q + 1, q + num_in, [](const Point<T>& A, const Point<T>& B) -> bool {
        T temp = cross_2d<T>(A, B);
        if (fabs(temp) < 1e-6) {
          return dot_2d<T>(A, A) < dot_2d<T>(B, B);
        } else {
          return temp > 0;
        }
      });
''' + ML_NMS_ROTATED_HEADER3

ML_NMS_ROTATED_CUDA_HEADER = ML_NMS_ROTATED_HEADER1 + r'''
#define HOST_DEVICE __host__ __device__
#define HOST_DEVICE_INLINE HOST_DEVICE __forceinline__
''' + ML_NMS_ROTATED_HEADER2 + r'''
  // CUDA version
  // In the future, we can potentially use thrust
  // for sorting here to improve speed (though not guaranteed)
  for (int i = 1; i < num_in - 1; i++) {
    for (int j = i + 1; j < num_in; j++) {
      T crossProduct = cross_2d<T>(q[i], q[j]);
      if ((crossProduct < -1e-6) ||
          (fabs(crossProduct) < 1e-6 && dist[i] > dist[j])) {
        auto q_tmp = q[i];
        q[i] = q[j];
        q[j] = q_tmp;
        auto dist_tmp = dist[i];
        dist[i] = dist[j];
        dist[j] = dist_tmp;
      }
    }
  }
''' + ML_NMS_ROTATED_HEADER3 + r'''
template <typename T>
__global__ void nms_rotated_cuda_kernel(
    const int n_boxes,
    const float iou_threshold,
    const T* dev_boxes,
    unsigned long long* dev_mask) {
  // nms_rotated_cuda_kernel is modified from torchvision's nms_cuda_kernel

  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

  // if (row_start > col_start) return;

  const int row_size =
      min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
      min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  // Compared to nms_cuda_kernel, where each box is represented with 4 values
  // (x1, y1, x2, y2), each rotated box is represented with 5 values
  // (x_center, y_center, width, height, angle_degrees) here.
  __shared__ T block_boxes[threadsPerBlock * BOX_LENGTH];
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x * BOX_LENGTH + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 0];
    block_boxes[threadIdx.x * BOX_LENGTH + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 1];
    block_boxes[threadIdx.x * BOX_LENGTH + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 2];
    block_boxes[threadIdx.x * BOX_LENGTH + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 3];
    block_boxes[threadIdx.x * BOX_LENGTH + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 4];
    if (BOX_LENGTH >= 6)
      block_boxes[threadIdx.x * BOX_LENGTH + 5] =
          dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * BOX_LENGTH + 5];
  }
  __syncthreads();

  if (threadIdx.x < row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const T* cur_box = dev_boxes + cur_box_idx * BOX_LENGTH;
    int i = 0;
    unsigned long long t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (i = start; i < col_size; i++) {
      // Instead of devIoU used by original horizontal nms, here
      // we use the single_box_iou_rotated function from box_iou_rotated_utils.h
      if (single_box_iou_rotated<T>(cur_box, block_boxes + i * BOX_LENGTH) >
          iou_threshold) {
        t |= 1ULL << i;
      }
    }
    const int col_blocks = CeilDIV(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t;
  }
}
'''

ML_NMS_ROTATED_CPU_SRC = r'''
  @alias(dets,in0)
  @alias(order_t,in1)
  @alias(suppressed_t,in2)
  @alias(keep_t,out0)

  const int ndets = dets_shape0;
  auto suppressed = suppressed_t_p;
  auto keep = keep_t_p;
  auto order = order_t_p;

  memset(keep,0,keep_t->size);

  int num_to_keep = 0;

  for (int _i = 0; _i < ndets; _i++) {
    auto i = order[_i];
    if (suppressed[i] == 1) {
      continue;
    }

    keep[i] = true;

    for (int _j = _i + 1; _j < ndets; _j++) {
      auto j = order[_j];
      if (suppressed[j] == 1) {
        continue;
      }

      auto ovr = single_box_iou_rotated(dets_p+i*BOX_LENGTH, dets_p+j*BOX_LENGTH);
      if (ovr >= iou_threshold) {
        suppressed[j] = 1;
      }
    }
  }
'''
ML_NMS_ROTATED_CUDA_SRC = r'''
  @alias(order,in0)
  @alias(dets_sorted,in1)
  @alias(keep,out0)

  cudaMemsetAsync(keep_p,0,keep->size);

  const int dets_num = dets_sorted_shape0;


  const int col_blocks = CeilDIV(dets_num, threadsPerBlock);

  int matrices_size = dets_num * col_blocks * sizeof(unsigned long long);
  size_t mask_allocation;
  unsigned long long* mask_p = (unsigned long long*)exe.allocator->alloc(matrices_size, mask_allocation);

  dim3 blocks(col_blocks, col_blocks);
  dim3 threads(threadsPerBlock);

  nms_rotated_cuda_kernel<<<blocks, threads, 0>>>(
            dets_num,
            iou_threshold,
            dets_sorted_p,
            mask_p);

  checkCudaErrors(cudaDeviceSynchronize());

  std::vector<unsigned long long> remv(col_blocks);
  memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);

  for (int i = 0; i < dets_num; i++) {
    int nblock = i / threadsPerBlock;
    int inblock = i % threadsPerBlock;

    if (!(remv[nblock] & (1ULL << inblock))) {
      keep_p[order_p[i]] = true;
      unsigned long long* p = mask_p + i * col_blocks;
      for (int j = nblock; j < col_blocks; j++) {
        remv[j] |= p[j];
      }
    }
  }
exe.allocator->free(mask_p, matrices_size, mask_allocation);
'''


def nms_rotated_cpu(dets, order_t, iou_threshold, box_length=6):
    suppressed_t = jt.zeros((dets.shape[0], )).uint8()

    keep = jt.code((dets.shape[0], ),
                   'bool', [dets, order_t, suppressed_t],
                   cpu_header=f'''
        #define BOX_LENGTH {box_length}
        ''' + ML_NMS_ROTATED_CPU_HEADER,
                   cpu_src=f'const float iou_threshold = {iou_threshold};' +
                   ML_NMS_ROTATED_CPU_SRC)
    return keep


def nms_rotated_cuda(dets, order_t, iou_threshold, box_length=6):
    dets_sorted = dets[order_t, :]
    keep = jt.code((dets.shape[0], ),
                   'bool', [order_t, dets_sorted],
                   cuda_header=f'''
        #define BOX_LENGTH {box_length}
        ''' + ML_NMS_ROTATED_CUDA_HEADER,
                   cuda_src=f'const float iou_threshold = {iou_threshold};' +
                   ML_NMS_ROTATED_CUDA_SRC)
    return keep


def ml_nms_rotated(dets, scores, labels, iou_threshold):
    assert dets.numel() > 0 and dets.ndim == 2
    assert dets.dtype == scores.dtype
    dets = jt.contrib.concat([dets, labels.cast(dets.dtype).unsqueeze(1)],
                             dim=1)
    order_t, _ = scores.argsort(0, descending=True)

    if jt.flags.use_cuda:
        keep = nms_rotated_cuda(dets, order_t, iou_threshold)
    else:
        keep = nms_rotated_cpu(dets, order_t, iou_threshold)
    return jt.where(keep)[0]


def nms_rotated(dets, scores, iou_threshold):
    if dets.numel() == 0:
        return jt.array([])
    assert dets.numel() > 0 and dets.ndim == 2
    assert dets.dtype == scores.dtype
    order_t, _ = scores.argsort(0, descending=True)

    if jt.flags.use_cuda:
        keep = nms_rotated_cuda(dets, order_t, iou_threshold, box_length=5)
    else:
        keep = nms_rotated_cpu(dets, order_t, iou_threshold, box_length=5)
    return jt.where(keep)[0]


def multiclass_nms_rotated(multi_bboxes,
                           multi_scores,
                           score_thr,
                           nms_cfg,
                           max_num=-1,
                           score_factors=None):
    """NMS for multi-class bboxes.

    Args:
        multi_bboxes (Tensor): shape (n, #class*5) or (n, 5)
        multi_scores (Tensor): shape (n, #class), where the 0th column
            contains scores of the background class, but this will be ignored.
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold
        max_num (int): if there are more than max_num bboxes after NMS,
            only top max_num will be kept.
        score_factors (Tensor): The factors multiplied to scores before
            applying NMS
    Returns:
        tuple: (bboxes, labels), tensors of shape (k, 6) and (k, 1). Labels
            are 0-based.
    """
    num_classes = multi_scores.size(1) - 1
    # exclude background category
    if multi_bboxes.shape[1] > 5:
        bboxes = multi_bboxes.view(multi_scores.size(0), -1, 5)[:, 1:]
    else:
        bboxes = multi_bboxes[:, None].expand(
            (multi_bboxes.shape[0], num_classes, 5))
    scores = multi_scores[:, 1:]

    # filter out boxes with low scores
    valid_mask = scores > score_thr
    bboxes = bboxes[valid_mask]
    if score_factors is not None:
        scores = scores * score_factors[:, None]
    scores = scores[valid_mask]
    labels = valid_mask.nonzero()[:, 1]

    if bboxes.numel() == 0:
        return jt.zeros((0, 6)), jt.zeros((0, )).int()
    nms_cfg_ = nms_cfg.copy()
    nms_type = nms_cfg_.pop('type', 'nms')
    iou_thr = nms_cfg_.pop('iou_thr', 0.1)
    keep = ml_nms_rotated(bboxes, scores, labels, iou_thr)
    bboxes = bboxes[keep]
    scores = scores[keep]
    labels = labels[keep]

    inds, _ = scores.argsort(descending=True)

    if keep.size(0) > max_num:
        inds = inds[:max_num]
    bboxes = bboxes[inds]
    scores = scores[inds]
    labels = labels[inds]

    return jt.contrib.concat([bboxes, scores[:, None]], 1), labels


def test_ml_nms_rotated():
    dets = jt.array([[0, 0, 1, 1, 0], [0, 0, 0.5, 0.5, 0.3],
                     [0, 0, 0.9, 0.9, 0]])
    scores = jt.array([0.1, 0.2, 0.3])
    labels = jt.array([1, 1, 1])
    print(ml_nms_rotated(dets, scores, labels, 0.3))
    print(nms_rotated(dets, scores, 0.3))


if __name__ == '__main__':
    test_ml_nms_rotated()

    jt.flags.use_cuda = 1
    test_ml_nms_rotated()

jittordet/ops/nms_rotated.py
JittorCV/jittordet
import math

import jittor as jt
from jittor import init, nn
from jittor.misc import _pair

__all__ = ['ORConv2d', 'RotationInvariantPooling']

ARF_CUDA_HEADER = r'''
typedef unsigned long uint64;
typedef unsigned int uint32;
typedef unsigned short uint16;
typedef unsigned char uint8;
#define CeilDIV(a,b) ((a+b-1)/b)
#define CUDA_KERNEL_LOOP(i, n)  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n;i += blockDim.x * gridDim.x)
template <typename Dtype>
__global__ void ARF_forward_cuda_kernel(
  const long nthreads,
  const Dtype* weight_data,
  const uint8* indices_data,
  const uint16 nInputPlane,
  const uint16 nOutputPlane,
  const uint8 nOrientation,
  const uint8 nRotation,
  const uint16 nEntry,
  Dtype* output_data)
{
  CUDA_KERNEL_LOOP(n, nthreads) {
    uint16 l = n % nEntry;
    uint16 j = (n / nEntry) % nInputPlane;
    uint16 i = n / nEntry / nInputPlane;
    uint8 k;
    Dtype val = *(weight_data + n);
    for (k = 0; k < nRotation; k++) {
      uint16 index = (uint16)(*(indices_data + l * nRotation + k)) - 1;
      Dtype *target = output_data + i * (nRotation * nInputPlane * nEntry)
                                  + k * (nInputPlane * nEntry)
                                  + j * (nEntry)
                                  + index;
      *target = val;
    }
  }
}
template <typename Dtype>
__global__ void ARF_backward_cuda_kernel(
  const long nthreads,
  const Dtype* gradWeight_data,
  const uint8* indices_data,
  const uint16 nInputPlane,
  const uint16 nOutputPlane,
  const uint8 nOrientation,
  const uint8 nRotation,
  const uint16 nEntry,
  Dtype* weight_data)
{
  CUDA_KERNEL_LOOP(n, nthreads) {
      uint16 l = n % nEntry;
      uint16 j = (n / nEntry) % nInputPlane;
      uint16 i = n / nEntry / nInputPlane;
      uint8 k;
      Dtype *val = weight_data + n;
      *val = 0;
      for (k = 0; k < nRotation; k++) {
          uint16 index = (uint16)(*(indices_data + l * nRotation + k)) - 1;
          Dtype target = *(gradWeight_data + i * (nRotation * nInputPlane * nEntry)
                                           + k * (nInputPlane * nEntry)
                                           + j * (nEntry)
                                           + index);
          *val = *val + target;
      }
  }
}
'''

ARF_CUDA_SRC = r'''
    @alias(weight,in0)
    @alias(indices,in1)
    @alias(output,out0)
    const uint16 nOutputPlane = weight_shape0;
    const uint16 nInputPlane = weight_shape1;
    const uint8 nOrientation = weight_shape2;
    const uint8 kH = weight_shape3;
    const uint8 kW = weight_shape4;
    const uint8 nRotation = indices_shape3;

    const uint16 nEntry = nOrientation * kH * kW;
    const long output_size = nOutputPlane * nInputPlane * nEntry;

    dim3 grid(std::min(CeilDIV(output_size, 512L), 4096L));
    dim3 block(512);

    ARF_forward_cuda_kernel<<<grid, block, 0>>>(
            output_size,
            weight_p,
            indices_p,
            nInputPlane,
            nOutputPlane,
            nOrientation,
            nRotation,
            nEntry,
            output_p);
'''
ARF_CUDA_GRAD_SRC = r'''
    @alias(indices,in0)
    @alias(gradOutput,in1)
    @alias(gradWeight,out0)
    const uint8 nOrientation = indices_shape0;
    const uint8 kH = indices_shape1;
    const uint8 kW = indices_shape2;
    const uint8 nRotation = indices_shape3;
    const uint16 nOutputPlane = gradOutput_shape0 / nRotation;
    const uint16 nInputPlane = gradOutput_shape1 / nOrientation;

    const uint16 nEntry = nOrientation * kH * kW;
    const long count = nOutputPlane * nInputPlane * nEntry;


    dim3 grid(std::min(CeilDIV(count, 512L), 4096L));
    dim3 block(512);

    ARF_backward_cuda_kernel<<<grid, block, 0>>>(
         count,
         gradOutput_p,
         indices_p,
         nInputPlane,
         nOutputPlane,
         nOrientation,
         nRotation,
         nEntry,
         gradWeight_p);
'''
ARF_CPU_HEADER = r'''
typedef unsigned long uint64;
typedef unsigned int uint32;
typedef unsigned short uint16;
typedef unsigned char uint8;
template <typename T>
void ARF_forward_cpu_kernel(
  const T* weightData,
  const uint8* indicesData,
  const uint16 nOutputPlane,
  const uint16 nInputPlane,
  const uint8 nOrientation,
  const uint8 kH,
  const uint8 kW,
  const uint8 nRotation,
  T* outputData)
{
  const uint16 nEntry = nOrientation * kH * kW;
  uint16 i, j, l;
  uint8 k;

#pragma omp parallel for private(i, j, l, k)
  for (i = 0; i < nOutputPlane; i++) {
    for (j = 0; j < nInputPlane; j++) {
      for (l = 0; l < nEntry; l++) {
        uint16 weightIndex = i * nInputPlane * nEntry
                             + j * nEntry
                             + l;
        T val = *(weightData + weightIndex);
        // T val = *(weightData++);
        for (k = 0; k < nRotation; k++) {
          uint16 index = (uint16)(*(indicesData + l * nRotation + k)) - 1;
          T *target = outputData + i * (nRotation * nInputPlane * nEntry)
                                 + k * (nInputPlane * nEntry)
                                 + j * (nEntry)
                                 + index;
          *target = val;
        }
      }
    }
  }
}
template <typename T>
void ARF_backward_cpu_kernel(
  const uint8* indicesData,
  const T* gradOutputData,
  const uint16 nOutputPlane,
  const uint16 nInputPlane,
  const uint8 nOrientation,
  const uint8 kH,
  const uint8 kW,
  const uint8 nRotation,
  T* gradInputData)
{
  const uint16 nEntry = nOrientation * kH * kW;
  uint16 i, j, l;
  uint8 k;

#pragma omp parallel for private(i, j, l, k)
  for (i = 0; i < nOutputPlane; i++) {
    for (j = 0; j < nInputPlane; j++) {
      for (l = 0; l < nEntry; l++) {
        uint16 gradInputIndex = i * nInputPlane * nEntry
                                + j * nEntry
                                + l;
        T *val = gradInputData + gradInputIndex;
        // T *val = gradInputData++;
        *val = 0;
        for (k = 0; k < nRotation; k++) {
          uint16 index = (uint16)(*(indicesData + l * nRotation + k)) - 1;
          const T *target = gradOutputData + i * (nRotation * nInputPlane * nEntry)
                                           + k * (nInputPlane * nEntry)
                                           + j * (nEntry)
                                           + index;
          *val = *val + *target;
        }
      }
    }
  }
}
'''
ARF_CPU_SRC = r'''
    @alias(weight,in0)
    @alias(indices,in1)
    @alias(output,out0)
    const uint16 nOutputPlane = weight_shape0;
    const uint16 nInputPlane = weight_shape1;
    const uint8 nOrientation = weight_shape2;
    const uint8 kH = weight_shape3;
    const uint8 kW = weight_shape4;
    const uint8 nRotation = indices_shape3;

    ARF_forward_cpu_kernel(
         weight_p,
         indices_p,
         nOutputPlane,
         nInputPlane,
         nOrientation,
         kH,
         kW,
         nRotation,
         output_p);
'''

ARF_CPU_GRAD_SRC = r'''
    @alias(indices,in0)
    @alias(gradOutput,in1)
    @alias(gradInput,out0)
    const uint8 nOrientation = indices_shape0;
    const uint8 kH = indices_shape1;
    const uint8 kW = indices_shape2;
    const uint8 nRotation = indices_shape3;
    const uint16 nOutputPlane = gradOutput_shape0 / nRotation;
    const uint16 nInputPlane = gradOutput_shape1 / nOrientation;

    ARF_backward_cpu_kernel(
         indices_p,
         gradOutput_p,
         nOutputPlane,
         nInputPlane,
         nOrientation,
         kH,
         kW,
         nRotation,
         gradInput_p);
'''


def arf_forward(input, indices):
    assert input.ndim == 5, 'only supports a batch of ARFs.'
    assert input.dtype == 'float32' and indices.dtype == 'uint8'
    nOutputPlane, nInputPlane, nOrientation, kH, kW = input.size()
    nRotation = indices.shape[3]

    output_shape = (nOutputPlane * nRotation, nInputPlane * nOrientation, kH,
                    kW)

    output = jt.code(
        output_shape,
        input.dtype, [input, indices],
        cpu_header=ARF_CPU_HEADER,
        cpu_src=ARF_CPU_SRC,
        cuda_header=ARF_CUDA_HEADER,
        cuda_src=ARF_CUDA_SRC)
    return output


def arf_backward(indices, grad_output):
    assert indices.ndim == 4 and indices.dtype == 'uint8' and grad_output.dtype == 'float32'

    nOrientation, kH, kW, nRotation = indices.size()
    nOutputPlane = grad_output.shape[0] // nRotation
    nInputPlane = grad_output.shape[1] // nOrientation

    output_shape = (nOutputPlane, nInputPlane, nOrientation, kH, kW)
    output = jt.code(
        output_shape,
        grad_output.dtype, [indices, grad_output],
        cpu_header=ARF_CPU_HEADER,
        cpu_src=ARF_CPU_GRAD_SRC,
        cuda_header=ARF_CUDA_HEADER,
        cuda_src=ARF_CUDA_GRAD_SRC)
    return output


RIE_CPU_HEADER = r'''
#define FLT_MAX 3.402823466e+38F
typedef unsigned long uint64;
typedef unsigned int uint32;
typedef unsigned short uint16;
typedef unsigned char uint8;

template <typename T>
void RIE_forward_cpu_kernel(
  const T* feature_data,
  uint8* mainDirection_data,
  T* aligned_data,
  const uint8 nOrientation,
  const uint16 nBatch,
  const uint16 nFeature)
{
  uint16 i;
  uint16 j;
  uint8 l;

  #pragma omp parallel for private(i, j, l)
  for (i = 0; i < nBatch; i++) {
    for (j = 0; j < nFeature; j++) {
      uint8 *direction = mainDirection_data + i * nFeature + j;
      T maxVal = -FLT_MAX;
      for (l = 0; l < nOrientation; l++) {
        T val = *(feature_data + i * (nFeature * nOrientation)
                               + j * (nOrientation)
                               + l);
        if (val > maxVal) {
          maxVal = val;
          *direction = l;
        }
      }
      for (l = 0; l < nOrientation; l++) {
        T src = *(feature_data + i * (nFeature * nOrientation)
                               + j * (nOrientation)
                               + l);
        uint8 alignedIndex = (l - (uint8)*direction + nOrientation) % nOrientation;
        T *target = aligned_data + i * (nFeature * nOrientation)
                                 + j * (nOrientation)
                                 + alignedIndex;
        *target = src;
      }
    }
  }
}

template <typename T>
void RIE_backward_cpu_kernel(
  const uint8* mainDirection_data,
  const T* gradOutput_data,
  const uint8 nOrientation,
  const uint16 nBatch,
  const uint16 nFeature,
  T* gradInput_data)
{
  uint16 i;
  uint16 j;
  uint8 l;

  #pragma omp parallel for private(i, j, l)
  for (i = 0; i < nBatch; i++) {
    for (j = 0; j < nFeature; j++) {
      uint8 direction = *(mainDirection_data + i * nFeature + j);
      for (l = 0; l < nOrientation; l++) {
        T src = *(gradOutput_data + i * (nFeature * nOrientation)
                                  + j * (nOrientation)
                                  + l);
        uint8 alignedIndex = (l + direction) % nOrientation;
        T *target = gradInput_data + i * (nFeature * nOrientation)
                                   + j * (nOrientation)
                                   + alignedIndex;
        *target = src;
      }
    }
  }
}
'''
RIE_CPU_SRC = r'''
    @alias(feature,in0)
    @alias(mainDirection,out0)
    @alias(aligned,out1)
    memset(aligned_p,0,aligned->size);
    const uint16 nBatch = feature_shape0;
    const uint16 nChannel = feature_shape1;
    const uint16 nFeature = nChannel / nOrientation;
    RIE_forward_cpu_kernel(
        feature_p,
        mainDirection_p,
        aligned_p,
        nOrientation,
        nBatch,
        nFeature);

'''
RIE_CPU_GRAD_SRC = r'''
    @alias(mainDirection,in0)
    @alias(gradOutput,in1)
    @alias(gradInput,out0)
    memset(gradInput_p,0,gradInput->size)

    const uint16 nBatch = mainDirection_shape0;
    const uint16 nFeature = mainDirection_shape1;

    RIE_backward_cpu_kernel(
         mainDirection_p,
         gradOutput_p,
         nOrientation,
         nBatch,
         nFeature,
         gradInput_p);
'''

RIE_CUDA_HEADER = r'''
#define FLT_MAX 3.402823466e+38F
typedef unsigned long uint64;
typedef unsigned int uint32;
typedef unsigned short uint16;
typedef unsigned char uint8;
#define CeilDIV(a,b) ((a+b-1)/b)
#define CUDA_KERNEL_LOOP(i, n)  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n;i += blockDim.x * gridDim.x)
template <typename Dtype>
__global__ void RIE_forward_cuda_kernel(
  const uint32 nthreads,
  const Dtype* feature_data,
  const uint16 nBatch,
  const uint16 nFeature,
  const uint8 nOrientation,
  uint8* mainDirection_data,
  Dtype* aligned_data)
{
  CUDA_KERNEL_LOOP(n, nthreads) {
    const uint16 j = n % nFeature;
    const uint16 i = n / nFeature;
    uint8 l;

    uint8 *direction = mainDirection_data + i * nFeature + j;
    Dtype maxVal = -FLT_MAX;
    for (l = 0; l < nOrientation; l++) {
      Dtype val = *(feature_data + i * (nFeature * nOrientation)
                                 + j * (nOrientation)
                                 + l);
      if (val > maxVal) {
        maxVal = val;
        *direction = l;
      }
    }
    for (l = 0; l < nOrientation; l++) {
      Dtype src = *(feature_data + i * (nFeature * nOrientation)
                                 + j * (nOrientation)
                                 + l);
      uint8 alignedIndex = ((l - (uint8)*direction) + nOrientation) % nOrientation;
      Dtype *target = aligned_data + i * (nFeature * nOrientation)
                                   + j * (nOrientation)
                                   + alignedIndex;
      *target = src;
    }
  }
}

template <typename Dtype>
__global__ void RIE_backward_cuda_kernel(
  const uint32 nthreads,
  const Dtype* aligned_data,
  const uint8* mainDirection_data,
  const uint16 nBatch,
  const uint16 nFeature,
  const uint8 nOrientation,
  Dtype* feature_data)
{
  CUDA_KERNEL_LOOP(n, nthreads) {
    uint8 l;
    const uint16 j = n % nFeature;
    const uint16 i = n / nFeature;
    const uint8 direction = *(mainDirection_data + i * nFeature + j);
    for (l = 0; l < nOrientation; l++) {
      Dtype src = *(aligned_data + i * (nFeature * nOrientation)
                                 + j * (nOrientation)
                                 + l);
      uint8 alignedIndex = (l + direction) % nOrientation;
      Dtype *target = feature_data + i * (nFeature * nOrientation)
                                   + j * (nOrientation)
                                   + alignedIndex;
      *target = src;
    }
  }
}
'''
RIE_CUDA_SRC = r'''
    @alias(feature,in0)
    @alias(mainDirection,out0)
    @alias(aligned,out1)
    const uint16 nBatch = feature_shape0;
    const uint16 nChannel = feature_shape1;
    const uint16 nFeature = nChannel / nOrientation;

    cudaMemsetAsync(aligned_p,0,aligned->size);
    const long count = nBatch * nFeature;
    dim3 grid(std::min(CeilDIV(count, 512L), 4096L));
    dim3 block(512);

    RIE_forward_cuda_kernel<<<grid, block, 0>>>(
            count,
            feature_p,
            nBatch,
            nFeature,
            nOrientation,
            mainDirection_p,
            aligned_p);
'''
RIE_CUDA_GRAD_SRC = r'''
    @alias(mainDirection,in0)
    @alias(gradOutput,in1)
    @alias(gradInput,out0)
    cudaMemsetAsync(gradInput_p,0,gradInput->size);

    const uint16 nBatch = mainDirection_shape0;
    const uint16 nFeature = mainDirection_shape1;
    const long count = nBatch * nFeature;

    dim3 grid(std::min(CeilDIV(count, 512L), 4096L));
    dim3 block(512);

    RIE_backward_cuda_kernel<<<grid, block, 0>>>(
            count,
            gradOutput_p,
            mainDirection_p,
            nBatch,
            nFeature,
            nOrientation,
            gradInput_p);
'''


def rie_forward(feature, nOrientation):
    assert feature.ndim == 4, 'only supports a batch of RIEs.'
    assert feature.size(2) == 1 and feature.size(
        3) == 1, 'mH x mW should be 1x1.'
    nBatch = feature.size(0)
    nChannel = feature.size(1)
    nFeature = nChannel // nOrientation

    prefix_src = f'const uint8 nOrientation = {nOrientation};'

    output = jt.code([(nBatch, nFeature), feature.shape],
                     ['uint8', feature.dtype], [feature],
                     cpu_header=RIE_CPU_HEADER,
                     cpu_src=prefix_src + RIE_CPU_SRC,
                     cuda_header=RIE_CUDA_HEADER,
                     cuda_src=prefix_src + RIE_CUDA_SRC)
    return output


def rie_backward(mainDirection, grad_output, nOrientation):
    prefix_src = f'const uint8 nOrientation = {nOrientation};'
    output = jt.code(
        grad_output.shape,
        grad_output.dtype, [mainDirection, grad_output],
        cpu_header=RIE_CPU_HEADER,
        cpu_src=prefix_src + RIE_CPU_GRAD_SRC,
        cuda_header=RIE_CUDA_HEADER,
        cuda_src=prefix_src + RIE_CUDA_GRAD_SRC)
    return output


class _ActiveRotatingFilter(jt.Function):

    def execute(self, input, indices):
        indices = indices.uint8()
        self.input = input
        output = arf_forward(input, indices)
        self.indices = indices
        return output

    def grad(self, grad_output):
        indices = self.indices
        input = self.input
        grad_input = arf_backward(indices, grad_output)
        return grad_input, None


class _RotationInvariantEncoding(jt.Function):

    def execute(self, input, nOrientation):
        self.nOrientation = nOrientation
        mainDirection, output = rie_forward(input, nOrientation)
        self.saved_tensors = (input, mainDirection)
        return output, mainDirection

    def grad(self, grad_output, g_tmp=None):
        input, mainDirection = self.saved_tensors
        grad_input = rie_backward(mainDirection, grad_output,
                                  self.nOrientation)
        return grad_input, None


rotation_invariant_encoding = _RotationInvariantEncoding.apply
active_rotating_filter = _ActiveRotatingFilter.apply


class ActiveRotatingFilter(nn.Module):

    def __init__(self, indices):
        super(ActiveRotatingFilter, self).__init__()
        self.indices = indices

    def execute(self, input):
        return active_rotating_filter(input, self.indices)


class RotationInvariantEncoding(nn.Module):

    def __init__(self, nOrientation, return_direction=False):
        super(RotationInvariantEncoding, self).__init__()
        self.nOrientation = nOrientation
        self.return_direction = return_direction

    def execute(self, input):
        output, d = rotation_invariant_encoding(input, self.nOrientation)
        if self.return_direction:
            return output, d
        else:
            return output


class RotationInvariantPooling(nn.Module):

    def __init__(self, nInputPlane, nOrientation=8):
        super(RotationInvariantPooling, self).__init__()
        self.nInputPlane = nInputPlane
        self.nOrientation = nOrientation
        #TODO remove this
        hiddent_dim = int(nInputPlane / nOrientation)
        self.conv = nn.Sequential(
            nn.Conv2d(hiddent_dim, nInputPlane, 1, 1),
            nn.BatchNorm2d(nInputPlane),
        )

    def execute(self, x):
        # TODO remove it
        self.conv.eval()
        # x: [N, c, 1, w]
        ## first, max_pooling along orientation.
        N, c, h, w = x.size()
        x = x.view(N, -1, self.nOrientation, h, w)
        x = x.max(dim=2, keepdims=False)  # [N, nInputPlane/nOrientation, 1, w]
        # MODIFIED
        # x = self.conv(x) # [N, nInputPlane, 1, w]
        return x


class ORConv2d(nn.Conv2d):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 arf_config=None,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias=True):
        self.nOrientation, self.nRotation = _pair(arf_config)
        assert (math.log(self.nOrientation) +
                1e-5) % math.log(2) < 1e-3, 'invalid nOrientation {}'.format(
                    self.nOrientation)
        assert (math.log(self.nRotation) +
                1e-5) % math.log(2) < 1e-3, 'invalid nRotation {}'.format(
                    self.nRotation)

        super(ORConv2d, self).__init__(in_channels, out_channels, kernel_size,
                                       stride, padding, dilation, groups, bias)
        self.indices = self.get_indices().stop_grad()

        self.weight = jt.zeros(
            (out_channels, in_channels, self.nOrientation, *self.kernel_size))
        if bias:
            self.bias = jt.zeros((out_channels * self.nRotation, ))

        self.reset_parameters()

    def reset_parameters(self):
        n = self.in_channels * self.nOrientation
        for k in self.kernel_size:
            n *= k
        init.gauss_(self.weight, 0, math.sqrt(2.0 / n))

    def get_indices(self, mode='fast'):
        kernel_indices = {
            1: {
                0: (1, ),
                45: (1, ),
                90: (1, ),
                135: (1, ),
                180: (1, ),
                225: (1, ),
                270: (1, ),
                315: (1, )
            },
            3: {
                0: (1, 2, 3, 4, 5, 6, 7, 8, 9),
                45: (2, 3, 6, 1, 5, 9, 4, 7, 8),
                90: (3, 6, 9, 2, 5, 8, 1, 4, 7),
                135: (6, 9, 8, 3, 5, 7, 2, 1, 4),
                180: (9, 8, 7, 6, 5, 4, 3, 2, 1),
                225: (8, 7, 4, 9, 5, 1, 6, 3, 2),
                270: (7, 4, 1, 8, 5, 2, 9, 6, 3),
                315: (4, 1, 2, 7, 5, 3, 8, 9, 6)
            }
        }
        delta_orientation = 360 / self.nOrientation
        delta_rotation = 360 / self.nRotation
        kH, kW = self.kernel_size
        indices = jt.zeros(
            (self.nOrientation * kH * kW, self.nRotation)).uint8()
        for i in range(0, self.nOrientation):
            for j in range(0, kH * kW):
                for k in range(0, self.nRotation):
                    angle = delta_rotation * k
                    layer = (i + math.floor(
                        angle / delta_orientation)) % self.nOrientation
                    kernel = kernel_indices[kW][angle][j]
                    indices[i * kH * kW + j, k] = int(layer * kH * kW + kernel)
        return indices.view(self.nOrientation, kH, kW, self.nRotation)

    def rotate_arf(self):
        return active_rotating_filter(self.weight, self.indices)

    def execute(self, input):
        return nn.conv2d(input, self.rotate_arf(), self.bias, self.stride,
                         self.padding, self.dilation, self.groups)

    def __repr__(self):
        arf_config = '[{}]'.format(self.nOrientation) \
        if self.nOrientation == self.nRotation \
        else '[{}-{}]'.format(self.nOrientation, self.nRotation)
        s = (
            '{name}({arf_config} {in_channels}, {out_channels}, kernel_size={kernel_size}'
            ', stride={stride}')
        if self.padding != (0, ) * len(self.padding):
            s += ', padding={padding}'
        if self.dilation != (1, ) * len(self.dilation):
            s += ', dilation={dilation}'
        if self.output_padding != (0, ) * len(self.output_padding):
            s += ', output_padding={output_padding}'
        if self.groups != 1:
            s += ', groups={groups}'
        if self.bias is None:
            s += ', bias=False'
        s += ')'
        return s.format(
            name=self.__class__.__name__,
            arf_config=arf_config,
            **self.__dict__)


def test_arf():
    jt.flags.use_cuda = 1
    import math
    out_channels = 4
    in_channels = 2
    nOrientation = 8
    nRotation = 8
    kernel_size = 3

    def get_indices(nOrientation, nRotation, kernel_size, mode='fast'):
        kernel_indices = {
            1: {
                0: (1, ),
                45: (1, ),
                90: (1, ),
                135: (1, ),
                180: (1, ),
                225: (1, ),
                270: (1, ),
                315: (1, )
            },
            3: {
                0: (1, 2, 3, 4, 5, 6, 7, 8, 9),
                45: (2, 3, 6, 1, 5, 9, 4, 7, 8),
                90: (3, 6, 9, 2, 5, 8, 1, 4, 7),
                135: (6, 9, 8, 3, 5, 7, 2, 1, 4),
                180: (9, 8, 7, 6, 5, 4, 3, 2, 1),
                225: (8, 7, 4, 9, 5, 1, 6, 3, 2),
                270: (7, 4, 1, 8, 5, 2, 9, 6, 3),
                315: (4, 1, 2, 7, 5, 3, 8, 9, 6)
            }
        }
        delta_orientation = 360 / nOrientation
        delta_rotation = 360 / nRotation
        kH, kW = kernel_size
        indices = jt.zeros((nOrientation * kH * kW, nRotation)).uint8()
        for i in range(0, nOrientation):
            for j in range(0, kH * kW):
                for k in range(0, nRotation):
                    angle = delta_rotation * k
                    layer = (i + math.floor(
                        angle / delta_orientation)) % nOrientation
                    kernel = kernel_indices[kW][angle][j]
                    indices[i * kH * kW + j, k] = int(layer * kH * kW + kernel)
        return indices.view(nOrientation, kH, kW, nRotation)

    input = jt.randn(
        (out_channels, in_channels, nOrientation, kernel_size, kernel_size))
    indices = get_indices(nOrientation, nRotation, (kernel_size, kernel_size))
    output = active_rotating_filter(input, indices)
    print(output.size())
    g1 = jt.grad(output.sum(), input)
    print(g1.mean())

    input = jt.randn((8, 16, 32, 32))
    orconv = ORConv2d(
        16, int(16 / 8), kernel_size=3, padding=1, arf_config=(1, 8))
    output = orconv(input)
    print(output.mean())
    g = jt.grad(output.sum(), input)
    print(g.mean())


def test_rie():
    nOrientation = 8
    input = jt.randn(2, 8, 1, 1)
    output, _ = rotation_invariant_encoding(input, nOrientation)

    g1 = jt.grad(output.sum(), input)
    print(g1.mean())


def test_rip():
    inst = RotationInvariantPooling(512, 8)
    input = jt.randn(8, 512, 1, 25)
    output = inst(input)
    print(output.size())


if __name__ == '__main__':
    test_arf()
    test_rie()
    test_rip()

jittordet/ops/orn.py
JittorCV/jittordet
import jittor as jt
from jittor import nn
from jittor.misc import _pair

__all__ = ['PSROIAlign']

CUDA_HEADER = r'''
#include <cmath>
#include <cstdio>
#include <climits>
#define CUDA_1D_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
using namespace std;
__device__ float bilinear_interpolate(const float* bottom_data,
    const int height, const int width,
    float y, float x,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    //empty
    return 0;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  int y_low = (int) y;
  int x_low = (int) x;
  int y_high;
  int x_high;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (float) y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (float) x_low;
  } else {
    x_high = x_low + 1;
  }
  float ly = y - y_low;
  float lx = x - x_low;
  float hy = 1. - ly, hx = 1. - lx;
  // do bilinear interpolation
  float v1 = bottom_data[y_low * width + x_low];
  float v2 = bottom_data[y_low * width + x_high];
  float v3 = bottom_data[y_high * width + x_low];
  float v4 = bottom_data[y_high * width + x_high];
  float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}
__device__ void bilinear_interpolate_gradient(
    const int height, const int width,
    float y, float x,
    float & w1, float & w2, float & w3, float & w4,
    int & x_low, int & x_high, int & y_low, int & y_high,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    //empty
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  y_low = (int) y;
  x_low = (int) x;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (float) y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (float) x_low;
  } else {
    x_high = x_low + 1;
  }
  float ly = y - y_low;
  float lx = x - x_low;
  float hy = 1. - ly, hx = 1. - lx;
  // reference in forward
  // float v1 = bottom_data[y_low * width + x_low];
  // float v2 = bottom_data[y_low * width + x_high];
  // float v3 = bottom_data[y_high * width + x_low];
  // float v4 = bottom_data[y_high * width + x_high];
  // float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  return;
}
'''

CUDA_SRC = r'''
__global__ void PSRoIAlignForward(@ARGS_DEF,const int nthreads, const float* bottom_data,
    const int channels,const int height, const int width,const int aligned_height, const int aligned_width,
    const float* bottom_rois, float* top_data) {
            @PRECALC
    // printf("%d\n", out_dim);
    const float spatial_scale = @in2(0);
    const float sampling_ratio = @in2(1);
    // const int out_dim = @in2(2);
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
        // NOfloatE Mingtao: copied from light-head
        // the order of dims is changed because the original implementation is in floatensorFlow
        int n = index;
        int pw = n % aligned_width;
        n /= aligned_width;
        int ph = n % aligned_height;
        n /= aligned_height;
        int ctop = n % out_dim;
        n /= out_dim;
        const float* offset_bottom_rois = bottom_rois + n * 5;
        int roi_batch_ind = offset_bottom_rois[0];
        // Do not using rounding; this implementation detail is critical
        auto roi_start_w = offset_bottom_rois[1] * spatial_scale;
        auto roi_start_h = offset_bottom_rois[2] * spatial_scale;
        auto roi_end_w = offset_bottom_rois[3] * spatial_scale;
        auto roi_end_h = offset_bottom_rois[4] * spatial_scale;
        // Force malformed ROIs to be 1x1
        auto roi_width = fmaxf(roi_end_w - roi_start_w, 1.f);
        auto roi_height = fmaxf(roi_end_h - roi_start_h, 1.f);
        auto bin_size_h = roi_height / aligned_height;
        auto bin_size_w = roi_width / aligned_width;
        // **************************************************************************************
        int c = ( ctop * aligned_width + ph ) * aligned_width + pw;
        // **************************************************************************************
        const float* offset_bottom_data =
            bottom_data + (roi_batch_ind * channels + c) * height * width; //channels?
        // We use roi_bin_grid to sample the grid and mimic integral
        int roi_bin_grid_h = (sampling_ratio > 0)
            ? sampling_ratio
            : ceil(roi_height / aligned_height); // e.g., = 2
        int roi_bin_grid_w =
            (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / aligned_width);
        // We do average (integral) pooling inside a bin
        const float count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
        float output_val = 0.;
        for (int iy = 0; iy < roi_bin_grid_h; iy++) // e.g., iy = 0, 1
        {
            const float y = roi_start_h + ph * bin_size_h +
                (iy + .5f) * bin_size_h / roi_bin_grid_h;  // e.g., 0.5, 1.5
            for (int ix = 0; ix < roi_bin_grid_w; ix++) {
                const float x = roi_start_w + pw * bin_size_w +
                    (ix + .5f) * bin_size_w / roi_bin_grid_w;
                float val = bilinear_interpolate(
                        offset_bottom_data, height, width, y, x, index);
                output_val += val;
            }
        }
        output_val /= count;
        top_data[index] = output_val;
    }
}
@alias(input,in0);
@alias(rois,in1);
@alias(output,out0);
auto num_rois = rois_shape0;
auto channels = input_shape1;
auto height = input_shape2;
auto width = input_shape3;
auto aligned_height = output_shape2;
auto aligned_width = output_shape3;
auto output_size = num_rois * aligned_height * aligned_width * out_dim;
const int total_count = in1_shape0 * out0_shape2 * out0_shape3 * in0_shape1;
const int thread_per_block = 512L;
const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
PSRoIAlignForward<<<block_count, thread_per_block>>>(@ARGS,output_size,input_p,channels,
height, width,aligned_height,aligned_width,rois_p,output_p);
'''

CUDA_GRAD_SRC = [
    r'''
__global__ void PSRoIAlignBackwardFeature(@ARGS_DEF,const int nthreads, const float* top_diff,
    const int num_rois,
    const int channels, const int height, const int width,
    const int aligned_height, const int aligned_width,
    float* bottom_diff,
    const float* bottom_rois) {
         @PRECALC

@alias(input,in0)
@alias(rois,in1)
@alias(grad_input,out0)
@alias(grad,dout)
    const float spatial_scale = @in2(0);
    const float sampling_ratio = @in2(1);
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
        // NOfloatE Mingtao: copied from light-head
        // the order of dims is changed because the original implementation is in floatensorFlow
        int n = index;
        int pw = n % aligned_width;
        n /= aligned_width;
        int ph = n % aligned_height;
        n /= aligned_height;
        int ctop = n % out_dim;
        n /= out_dim;
        const float* offset_bottom_rois = bottom_rois + n * 5;
        int roi_batch_ind = offset_bottom_rois[0];
        // Do not using rounding; this implementation detail is critical
        float roi_start_w = offset_bottom_rois[1] * spatial_scale;
        float roi_start_h = offset_bottom_rois[2] * spatial_scale;
        float roi_end_w = offset_bottom_rois[3] * spatial_scale;
        float roi_end_h = offset_bottom_rois[4] * spatial_scale;
        // Force malformed ROIs to be 1x1
        float roi_width = fmaxf(roi_end_w - roi_start_w, 1.f);
        float roi_height = fmaxf(roi_end_h - roi_start_h, 1.f);
        float bin_size_h = roi_height / aligned_height;
        float bin_size_w = roi_width / aligned_width;
        int c = ( ctop * aligned_width + ph ) * aligned_width + pw;
        float* offset_bottom_diff = bottom_diff + (roi_batch_ind * channels + c) * height * width;

        int top_offset = (n * out_dim + ctop) * aligned_height * aligned_width;
        const float* offset_top_diff = top_diff + top_offset;
        const float top_diff_this_bin = offset_top_diff[ph * aligned_width + pw];
        // We use roi_bin_grid to sample the grid and mimic integral
        int roi_bin_grid_h = (sampling_ratio > 0)
            ? sampling_ratio
            : ceil(roi_height / aligned_height); // e.g., = 2
        int roi_bin_grid_w =
            (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / aligned_width);
        // We do average (integral) pooling inside a bin
        const float count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
        for (int iy = 0; iy < roi_bin_grid_h; iy++) // e.g., iy = 0, 1
        {
            const float y = roi_start_h + ph * bin_size_h +
                (iy + .5f) * bin_size_h / roi_bin_grid_h; // e.g., 0.5, 1.5
            for (int ix = 0; ix < roi_bin_grid_w; ix++) {
                const float x = roi_start_w + pw * bin_size_w +
                    (ix + .5f) * bin_size_w / roi_bin_grid_w;
                float w1, w2, w3, w4;
                int x_low, x_high, y_low, y_high;
                bilinear_interpolate_gradient(
                        height, width, y, x, w1, w2, w3, w4,
                        x_low, x_high, y_low, y_high, index);
                float g1 = top_diff_this_bin * w1 / count;
                float g2 = top_diff_this_bin * w2 / count;
                float g3 = top_diff_this_bin * w3 / count;
                float g4 = top_diff_this_bin * w4 / count;
                if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
                    atomicAdd(offset_bottom_diff + y_low * width + x_low, static_cast<float>(g1));
                    atomicAdd(offset_bottom_diff + y_low * width + x_high, static_cast<float>(g2));
                    atomicAdd(offset_bottom_diff + y_high * width + x_low, static_cast<float>(g3));
                    atomicAdd(offset_bottom_diff + y_high * width + x_high, static_cast<float>(g4));
                } // if
            } // ix
        } // iy
    } // CUDA_1D_KERNEL_LOOP
} // RoIAlignBackward
auto num_rois = rois_shape0;
auto channels = input_shape1;
auto height = input_shape2;
auto width = input_shape3;
auto aligned_height = grad_shape2;
auto aligned_width = grad_shape3;
auto output_size = num_rois * aligned_height * aligned_width * out_dim;
cudaMemsetAsync(grad_input_p,0,grad_input->size);
const int total_count = rois_shape0 * grad_shape2 * grad_shape3 * input_shape1;
const int thread_per_block = 512;
const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
PSRoIAlignBackwardFeature<<<block_count, thread_per_block>>>(@ARGS,output_size,grad_p,num_rois,
channels, height, width, aligned_height, aligned_width, grad_input_p, rois_p);
''', '', ''
]


def psroi_align(input, rois, output_size, spatial_scale, sampling_ratio,
                out_dim):
    output_size = _pair(output_size)
    options = jt.array([spatial_scale, sampling_ratio, out_dim])
    output_shapes = (rois.shape[0], input.shape[1], output_size[0],
                     output_size[1])
    inputs = [input, rois, options]
    output_types = input.dtype
    if rois.shape[0] == 0:
        return jt.zeros(output_shapes, input.dtype)
    output = jt.code(
        output_shapes,
        output_types,
        inputs,
        cuda_header=CUDA_HEADER,
        cuda_src=CUDA_SRC,
        cuda_grad_src=CUDA_GRAD_SRC)
    return output


class PSROIAlign(nn.Module):

    def __init__(self, output_size, spatial_scale, sampling_ratio, out_dim):
        super(PSROIAlign, self).__init__()
        self.output_size = output_size
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        self.out_dim = out_dim

    def execute(self, input, rois):
        return psroi_align(input, rois, self.output_size, self.spatial_scale,
                           self.sampling_ratio, self.out_dim)

    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 += ')'
        return tmpstr


def test_psroi_align():
    jt.flags.use_cuda = 1
    psroialign = PSROIAlign(7, 1.0, sampling_ratio=0, out_dim=23)

    feature = jt.rand((8, 512, 64, 64), dtype='float32')
    boxes = jt.array([[0, 1.2, 10, 3, 10], [3, 1, 67, 2, 34]])
    output = psroialign(feature, boxes)
    print(output.shape)
    # print(jt.grad(output,feature))


if __name__ == '__main__':
    test_psroi_align()

jittordet/ops/psroi_align.py
JittorCV/jittordet
import jittor as jt
from jittor import nn
from jittor.misc import _pair

__all__ = ['ROIAlign']

CUDA_HEADER = r'''
#include <cmath>
#include <cstdio>
#include <climits>
#define CUDA_1D_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
using namespace std;
__device__ float bilinear_interpolate(const float* bottom_data,
    const int height, const int width,
    float y, float x,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    //empty
    return 0;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  int y_low = (int) y;
  int x_low = (int) x;
  int y_high;
  int x_high;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (float) y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (float) x_low;
  } else {
    x_high = x_low + 1;
  }
  float ly = y - y_low;
  float lx = x - x_low;
  float hy = 1. - ly, hx = 1. - lx;
  // do bilinear interpolation
  float v1 = bottom_data[y_low * width + x_low];
  float v2 = bottom_data[y_low * width + x_high];
  float v3 = bottom_data[y_high * width + x_low];
  float v4 = bottom_data[y_high * width + x_high];
  float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}
__device__ void bilinear_interpolate_gradient(
    const int height, const int width,
    float y, float x,
    float & w1, float & w2, float & w3, float & w4,
    int & x_low, int & x_high, int & y_low, int & y_high,
    const int index /* index for debug only*/) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    //empty
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  y_low = (int) y;
  x_low = (int) x;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (float) y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (float) x_low;
  } else {
    x_high = x_low + 1;
  }
  float ly = y - y_low;
  float lx = x - x_low;
  float hy = 1. - ly, hx = 1. - lx;
  // reference in forward
  // float v1 = bottom_data[y_low * width + x_low];
  // float v2 = bottom_data[y_low * width + x_high];
  // float v3 = bottom_data[y_high * width + x_low];
  // float v4 = bottom_data[y_high * width + x_high];
  // float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  return;
}
__global__ void RoIAlignForward(const int nthreads, const float* bottom_data,
    const int channels,const int height, const int width,const int pooled_height, const int pooled_width,
    const float* bottom_rois, float* top_data,const float spatial_scale,const float  sampling_ratio) {
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    const float* offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    // Do not using rounding; this implementation detail is critical
    auto roi_start_w = offset_bottom_rois[1] * spatial_scale;
    auto roi_start_h = offset_bottom_rois[2] * spatial_scale;
    #if ROI_ALIGN_VERSION == 1
    auto roi_end_w = (offset_bottom_rois[3] + 1) * spatial_scale;
    auto roi_end_h = (offset_bottom_rois[4] + 1) * spatial_scale;
    auto roi_width = fmaxf(roi_end_w - roi_start_w, 0.);
    auto roi_height = fmaxf(roi_end_h - roi_start_h, 0.);
    #else
    auto roi_end_w = offset_bottom_rois[3] * spatial_scale;
    auto roi_end_h = offset_bottom_rois[4] * spatial_scale;
    auto roi_width = fmaxf(roi_end_w - roi_start_w, 1.);
    auto roi_height = fmaxf(roi_end_h - roi_start_h, 1.);
    #endif
    auto bin_size_h = static_cast<float>(roi_height) / static_cast<float>(pooled_height);
    auto bin_size_w = static_cast<float>(roi_width) / static_cast<float>(pooled_width);
    const float* offset_bottom_data = bottom_data + (roi_batch_ind * channels + c) * height * width;
    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sampling_ratio > 0) ? sampling_ratio : ceil(roi_height / pooled_height); // e.g., = 2
    int roi_bin_grid_w = (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width);
    // We do average (integral) pooling inside a bin
    const float count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
    float output_val = 0.;
    for (int iy = 0; iy < roi_bin_grid_h; iy ++) // e.g., iy = 0, 1
    {
      const float y = roi_start_h + ph * bin_size_h + static_cast<float>(iy + .5f) * bin_size_h / static_cast<float>(roi_bin_grid_h); // e.g., 0.5, 1.5
      for (int ix = 0; ix < roi_bin_grid_w; ix ++)
      {
        const float x = roi_start_w + pw * bin_size_w + static_cast<float>(ix + .5f) * bin_size_w / static_cast<float>(roi_bin_grid_w);
        float val = bilinear_interpolate(offset_bottom_data, height, width, y, x, index);
        output_val += val;
      }
    }
    output_val /= count;
    top_data[index] = output_val;
  }
}
__global__ void RoIAlignBackwardFeature(const int nthreads, const float* top_diff,const int num_rois,
       const int channels, const int height, const int width,const int pooled_height, const int pooled_width,
       float* bottom_diff,const float* bottom_rois,const float spatial_scale,const float  sampling_ratio) {
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    const float* offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    // Do not using rounding; this implementation detail is critical
    float roi_start_w = offset_bottom_rois[1] * spatial_scale;
    float roi_start_h = offset_bottom_rois[2] * spatial_scale;
    #if ROI_ALIGN_VERSION == 1
    float roi_end_w = (offset_bottom_rois[3] + 1) * spatial_scale;
    float roi_end_h = (offset_bottom_rois[4] + 1) * spatial_scale;
    float roi_width = fmaxf(roi_end_w - roi_start_w, (float)0.);
    float roi_height = fmaxf(roi_end_h - roi_start_h, (float)0.);
    #else
    float roi_end_w = offset_bottom_rois[3] * spatial_scale;
    float roi_end_h = offset_bottom_rois[4] * spatial_scale;
    float roi_width = fmaxf(roi_end_w - roi_start_w, (float)1.);
    float roi_height = fmaxf(roi_end_h - roi_start_h, (float)1.);
    #endif
    // float roi_start_w = round(offset_bottom_rois[1] * spatial_scale);
    // float roi_start_h = round(offset_bottom_rois[2] * spatial_scale);
    // float roi_end_w = round(offset_bottom_rois[3] * spatial_scale);
    // float roi_end_h = round(offset_bottom_rois[4] * spatial_scale);
    float bin_size_h = static_cast<float>(roi_height) / static_cast<float>(pooled_height);
    float bin_size_w = static_cast<float>(roi_width) / static_cast<float>(pooled_width);
    float* offset_bottom_diff = bottom_diff + (roi_batch_ind * channels + c) * height * width;
    int top_offset    = (n * channels + c) * pooled_height * pooled_width;
    const float* offset_top_diff = top_diff + top_offset;
    const float top_diff_this_bin = offset_top_diff[ph * pooled_width + pw];
    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sampling_ratio > 0) ? sampling_ratio : ceil(roi_height / pooled_height); // e.g., = 2
    int roi_bin_grid_w = (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width);
    // We do average (integral) pooling inside a bin
    const float count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
    for (int iy = 0; iy < roi_bin_grid_h; iy ++) // e.g., iy = 0, 1
    {
      const float y = roi_start_h + ph * bin_size_h + static_cast<float>(iy + .5f) * bin_size_h / static_cast<float>(roi_bin_grid_h); // e.g., 0.5, 1.5
      for (int ix = 0; ix < roi_bin_grid_w; ix ++)
      {
        const float x = roi_start_w + pw * bin_size_w + static_cast<float>(ix + .5f) * bin_size_w / static_cast<float>(roi_bin_grid_w);
        float w1, w2, w3, w4;
        int x_low, x_high, y_low, y_high;
        bilinear_interpolate_gradient(height, width, y, x,
            w1, w2, w3, w4,
            x_low, x_high, y_low, y_high,
            index);
        float g1 = top_diff_this_bin * w1 / count;
        float g2 = top_diff_this_bin * w2 / count;
        float g3 = top_diff_this_bin * w3 / count;
        float g4 = top_diff_this_bin * w4 / count;
        if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0)
        {
          atomicAdd(offset_bottom_diff + y_low * width + x_low, static_cast<float>(g1));
          atomicAdd(offset_bottom_diff + y_low * width + x_high, static_cast<float>(g2));
          atomicAdd(offset_bottom_diff + y_high * width + x_low, static_cast<float>(g3));
          atomicAdd(offset_bottom_diff + y_high * width + x_high, static_cast<float>(g4));
        } // if
      } // ix
    } // iy
  } // CUDA_1D_KERNEL_LOOP
} // RoIAlignBackward
'''


class _ROIAlign(jt.Function):

    def execute(self,
                input,
                rois,
                output_size,
                spatial_scale,
                sampling_ratio,
                version=0):
        self.input = input
        self.rois = rois
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        output_shapes = (rois.shape[0], input.shape[1], output_size[0],
                         output_size[1])
        self.version = version
        return jt.code(
            output_shapes,
            input.dtype, [input, rois],
            cuda_header=f"""
              #define ROI_ALIGN_VERSION {self.version}
              {CUDA_HEADER}""",
            cuda_src=f'''
              @alias(input,in0);
              @alias(rois,in1);
              @alias(output,out0);
              const float spatial_scale = {spatial_scale};
              const float  sampling_ratio = {sampling_ratio};
              auto num_rois = rois_shape0;
              auto channels = input_shape1;
              auto height = input_shape2;
              auto width = input_shape3;
              auto pooled_height = output_shape2;
              auto pooled_width = output_shape3;
              auto output_size = num_rois * pooled_height * pooled_width * channels;
              const int total_count = in1_shape0 * out0_shape2 * out0_shape3 * in0_shape1;
              const int thread_per_block = 512L;
              const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
              RoIAlignForward<<<block_count, thread_per_block>>>(output_size,input_p,channels,height, width,pooled_height,pooled_width,rois_p,output_p,spatial_scale,sampling_ratio);
              ''')

    def grad(self, output_grad):
        input, rois = self.input, self.rois
        input_grad = jt.code(
            input.shape,
            input.dtype, [input, rois, output_grad],
            cuda_header=f"""
                        #define ROI_ALIGN_VERSION {self.version}
                        {CUDA_HEADER}""",
            cuda_src=f'''
                        @alias(input,in0)
                        @alias(rois,in1)
                        @alias(grad,in2)
                        @alias(grad_input,out0)
                        const float spatial_scale = {self.spatial_scale};
                        const float  sampling_ratio = {self.sampling_ratio};
                        auto num_rois = rois_shape0;
                        auto channels = input_shape1;
                        auto height = input_shape2;
                        auto width = input_shape3;
                        auto pooled_height = grad_shape2;
                        auto pooled_width = grad_shape3;
                        auto output_size = num_rois * pooled_height * pooled_width * channels;
                        cudaMemsetAsync(grad_input_p,0,grad_input->size);
                        const int total_count = rois_shape0 * grad_shape2 * grad_shape3 * input_shape1;
                        const int thread_per_block = 512;
                        const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
                        RoIAlignBackwardFeature<<<block_count, thread_per_block>>>(output_size,grad_p,num_rois,channels,height,width,pooled_height,pooled_width,grad_input_p,rois_p,spatial_scale,sampling_ratio);
                        ''')
        return input_grad, None


roi_align = _ROIAlign.apply


class ROIAlign(nn.Module):

    def __init__(self,
                 output_size,
                 spatial_scale,
                 sampling_ratio=0,
                 version=0):
        super(ROIAlign, self).__init__()
        self.output_size = _pair(output_size)
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        self.version = version

    def execute(self, input, rois):
        res = roi_align(input, rois, self.output_size, self.spatial_scale,
                        self.sampling_ratio, self.version)
        return res

    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 += ', version=' + str(self.version)
        tmpstr += ')'
        return tmpstr


def test_roialign():
    jt.flags.use_cuda = 1
    roi_align = ROIAlign((7, 7), 1 / 16.)
    feature = jt.randn((2, 1024, 64, 64))
    roi = jt.array([[0, 20, 120, 80, 195.5], [1, 23, 56, 200, 300.5]])
    output = roi_align(feature, roi)
    print(output.shape, output.sum())
    print(jt.grad(output.exp(), feature))


if __name__ == '__main__':
    test_roialign()

jittordet/ops/roi_align.py
JittorCV/jittordet
import jittor as jt
from jittor import nn
from jittor.misc import _pair

__all__ = ['ROIAlign']

CUDA_HEADER = r'''
#include <cmath>
#include <cstdio>
#include <climits>
#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)
#define THREADS_PER_BLOCK 1024
inline int GET_BLOCKS(const int N) {
    int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    int max_block_num = 65000;
    return min(optimal_block_num, max_block_num);
}
template <typename scalar_t>
__device__ scalar_t bilinear_interpolate(const scalar_t *bottom_data,
                                         const int height, const int width,
                                         scalar_t y, scalar_t x) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    return 0;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  int y_low = (int)y;
  int x_low = (int)x;
  int y_high;
  int x_high;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (scalar_t)y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (scalar_t)x_low;
  } else {
    x_high = x_low + 1;
  }
  scalar_t ly = y - y_low;
  scalar_t lx = x - x_low;
  scalar_t hy = 1. - ly;
  scalar_t hx = 1. - lx;
  // do bilinear interpolation
  scalar_t lt = bottom_data[y_low * width + x_low];
  scalar_t rt = bottom_data[y_low * width + x_high];
  scalar_t lb = bottom_data[y_high * width + x_low];
  scalar_t rb = bottom_data[y_high * width + x_high];
  scalar_t w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  scalar_t val = (w1 * lt + w2 * rt + w3 * lb + w4 * rb);
  return val;
}
template <typename scalar_t>
__global__ void ROIAlignRotatedForward(const int nthreads, const scalar_t *bottom_data,
                                const scalar_t *bottom_rois,
                                const scalar_t spatial_scale,
                                const int sample_num, const int channels,
                                const int height, const int width,
                                const int pooled_height, const int pooled_width,
                                scalar_t *top_data) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    const scalar_t* offset_bottom_rois = bottom_rois + n * 6;
    int roi_batch_ind = offset_bottom_rois[0];
    // Do not using rounding; this implementation detail is critical
    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale;
    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale;
    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
    scalar_t theta = offset_bottom_rois[5];
    // Force malformed ROIs to be 1x1
    roi_width = max(roi_width, (scalar_t)1.);
    roi_height = max(roi_height, (scalar_t)1.);
    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) / static_cast<scalar_t>(pooled_height);
    scalar_t bin_size_w = static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);
    const scalar_t* offset_bottom_data =
        bottom_data + (roi_batch_ind * channels + c) * height * width;
    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sample_num > 0)
        ? sample_num
        : ceil(roi_height / pooled_height);  // e.g., = 2
    int roi_bin_grid_w =
        (sample_num > 0) ? sample_num : ceil(roi_width / pooled_width);

    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
    // Appropriate translation needs to be applied after.
    scalar_t roi_start_h = -roi_height / 2.0;
    scalar_t roi_start_w = -roi_width / 2.0;
    scalar_t cosscalar_theta = cos(theta);
    scalar_t sinscalar_theta = sin(theta);
    // We do average (integral) pooling inside a bin
    const scalar_t count = roi_bin_grid_h * roi_bin_grid_w;  // e.g. = 4
    scalar_t output_val = 0.;
    for (int iy = 0; iy < roi_bin_grid_h; iy++) {  // e.g., iy = 0, 1
        const scalar_t yy = roi_start_h + ph * bin_size_h +
            static_cast<scalar_t>(iy + .5f) * bin_size_h /
                static_cast<scalar_t>(roi_bin_grid_h);  // e.g., 0.5, 1.5
        for (int ix = 0; ix < roi_bin_grid_w; ix++) {
        const scalar_t xx = roi_start_w + pw * bin_size_w +
            static_cast<scalar_t>(ix + .5f) * bin_size_w /
                static_cast<scalar_t>(roi_bin_grid_w);
        // Rotate by theta around the center and translate
        // scalar_t x = xx * cosscalar_theta + yy * sinscalar_theta + roi_center_w;
        // scalar_t y = yy * cosscalar_theta - xx * sinscalar_theta + roi_center_h;
        scalar_t x = xx * cosscalar_theta - yy * sinscalar_theta + roi_center_w;
        scalar_t y = xx * sinscalar_theta + yy * cosscalar_theta + roi_center_h;
        scalar_t val = bilinear_interpolate<scalar_t>(
            offset_bottom_data, height, width, y, x);
        output_val += val;
        }
    }
    output_val /= count;
    top_data[index] = output_val;
    }
}
template <typename scalar_t>
__device__ void bilinear_interpolate_gradient(const int height, const int width,
                                              scalar_t y, scalar_t x,
                                              scalar_t &w1, scalar_t &w2,
                                              scalar_t &w3, scalar_t &w4,
                                              int &x_low, int &x_high,
                                              int &y_low, int &y_high) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }
  if (y <= 0) y = 0;
  if (x <= 0) x = 0;
  y_low = (int)y;
  x_low = (int)x;
  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (scalar_t)y_low;
  } else {
    y_high = y_low + 1;
  }
  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (scalar_t)x_low;
  } else {
    x_high = x_low + 1;
  }
  scalar_t ly = y - y_low;
  scalar_t lx = x - x_low;
  scalar_t hy = 1. - ly;
  scalar_t hx = 1. - lx;
  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
  return;
}
template <typename scalar_t>
__global__ void ROIAlignBackward(
    const int nthreads, const scalar_t *top_diff, const scalar_t *bottom_rois,
    const scalar_t spatial_scale, const int sample_num, const int channels,
    const int height, const int width, const int pooled_height,
    const int pooled_width, scalar_t *bottom_diff) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    const scalar_t* offset_bottom_rois = bottom_rois + n * 6;
    int roi_batch_ind = offset_bottom_rois[0];
    // Do not round
    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale;
    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale;
    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
    scalar_t theta = offset_bottom_rois[5];

    // Force malformed ROIs to be 1x1
    roi_width = max(roi_width, (scalar_t)1.);
    roi_height = max(roi_height, (scalar_t)1.);
    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) / static_cast<scalar_t>(pooled_height);
    scalar_t bin_size_w = static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);
    scalar_t* offset_bottom_diff =
        bottom_diff + (roi_batch_ind * channels + c) * height * width;
    int top_offset = (n * channels + c) * pooled_height * pooled_width;
    const scalar_t* offset_top_diff = top_diff + top_offset;
    const scalar_t top_diff_this_bin = offset_top_diff[ph * pooled_width + pw];
    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sample_num > 0)
        ? sample_num
        : ceil(roi_height / pooled_height);  // e.g., = 2
    int roi_bin_grid_w =
        (sample_num > 0) ? sample_num : ceil(roi_width / pooled_width);
    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
    // Appropriate translation needs to be applied after.
    scalar_t roi_start_h = -roi_height / 2.0;
    scalar_t roi_start_w = -roi_width / 2.0;
    scalar_t cosTheta = cos(theta);
    scalar_t sinTheta = sin(theta);
    // We do average (integral) pooling inside a bin
    const scalar_t count = roi_bin_grid_h * roi_bin_grid_w;  // e.g. = 4
    for (int iy = 0; iy < roi_bin_grid_h; iy++) {  // e.g., iy = 0, 1
      const scalar_t yy = roi_start_h + ph * bin_size_h +
          static_cast<scalar_t>(iy + .5f) * bin_size_h /
              static_cast<scalar_t>(roi_bin_grid_h);  // e.g., 0.5, 1.5
      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
        const scalar_t xx = roi_start_w + pw * bin_size_w +
            static_cast<scalar_t>(ix + .5f) * bin_size_w /
                static_cast<scalar_t>(roi_bin_grid_w);
        // Rotate by theta around the center and translate
        // scalar_t x = xx * cosTheta + yy * sinTheta + roi_center_w;
        // T y = yy * cosTheta - xx * sinTheta + roi_center_h;
        scalar_t x = xx * cosTheta - yy * sinTheta + roi_center_w;
        scalar_t y = xx * sinTheta + yy * cosTheta + roi_center_h;
        scalar_t w1, w2, w3, w4;
        int x_low, x_high, y_low, y_high;
        bilinear_interpolate_gradient<scalar_t>(
            height,
            width,
            y,
            x,
            w1,
            w2,
            w3,
            w4,
            x_low,
            x_high,
            y_low,
            y_high);
        scalar_t g1 = top_diff_this_bin * w1 / count;
        scalar_t g2 = top_diff_this_bin * w2 / count;
        scalar_t g3 = top_diff_this_bin * w3 / count;
        scalar_t g4 = top_diff_this_bin * w4 / count;
        if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
          atomicAdd(
              offset_bottom_diff + y_low * width + x_low, g1);
          atomicAdd(
              offset_bottom_diff + y_low * width + x_high, g2);
          atomicAdd(
              offset_bottom_diff + y_high * width + x_low, g3);
          atomicAdd(
              offset_bottom_diff + y_high * width + x_high, g4);
        }  // if
      }  // ix
    }  // iy
  }  // CUDA_1D_KERNEL_LOOP
}  // RoIAlignBackward
'''


class _RotatedROIAlign(jt.Function):

    def execute(self, input, rois, output_size, spatial_scale, sampling_ratio):
        self.input = input
        self.rois = rois
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        assert rois.shape[1] == 6
        output_shapes = (rois.shape[0], input.shape[1], output_size[0],
                         output_size[1])
        return jt.code(
            output_shapes,
            input.dtype, [input, rois],
            cuda_header=CUDA_HEADER,
            cuda_src=f'''
                            @alias(input,in0);
                            @alias(rois,in1);
                            @alias(output,out0);
                            const float spatial_scale = {spatial_scale};
                            const float  sampling_ratio = {sampling_ratio};
                            auto num_rois = rois_shape0;
                            auto channels = input_shape1;
                            auto height = input_shape2;
                            auto width = input_shape3;
                            auto pooled_height = output_shape2;
                            auto pooled_width = output_shape3;
                            auto output_size = num_rois * pooled_height * pooled_width * channels;
                            ROIAlignRotatedForward<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>> (
                                        output_size, input_p, rois_p, spatial_scale,
                                        sampling_ratio, channels, height, width, pooled_height,
                                        pooled_width, output_p);
                            ''')

    def grad(self, output_grad):
        input, rois = self.input, self.rois
        input_grad = jt.code(
            input.shape,
            input.dtype, [input, rois, output_grad],
            cuda_header=CUDA_HEADER,
            cuda_src=f'''
                            @alias(input,in0)
                            @alias(rois,in1)
                            @alias(grad,in2)
                            @alias(grad_input,out0)
                            const float spatial_scale = {self.spatial_scale};
                            const float  sampling_ratio = {self.sampling_ratio};
                            auto num_rois = rois_shape0;
                            auto channels = input_shape1;
                            auto height = input_shape2;
                            auto width = input_shape3;
                            auto pooled_height = grad_shape2;
                            auto pooled_width = grad_shape3;
                            auto output_size = num_rois * pooled_height * pooled_width * channels;
                            cudaMemsetAsync(grad_input_p,0,grad_input->size);
                             ROIAlignBackward<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>>(
                                    output_size, grad_p, rois_p, spatial_scale, sampling_ratio,
                                    channels, height, width, pooled_height, pooled_width,
                                    grad_input_p);
                            ''')
        return input_grad, None


roi_align = _RotatedROIAlign.apply


class ROIAlignRotated(nn.Module):

    def __init__(self, output_size, spatial_scale, sampling_ratio=0):
        super(ROIAlignRotated, self).__init__()
        self.output_size = _pair(output_size)
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio

    def execute(self, input, rois):
        return roi_align(input, rois, self.output_size, self.spatial_scale,
                         self.sampling_ratio)

    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 += ')'
        return tmpstr


def test_roialign():
    jt.flags.use_cuda = 1
    roi_align = ROIAlignRotated((7, 7), 1 / 16.)
    feature = jt.randn((2, 1024, 64, 64))
    roi = jt.array([[0, 20, 120, 80, 195.5, 0.3], [1, 23, 56, 200, 300.5,
                                                   0.2]])
    output = roi_align(feature, roi)
    print(output.shape, output.sum())
    print(jt.grad(output.exp(), feature))


if __name__ == '__main__':
    test_roialign()

jittordet/ops/roi_align_rotated.py
JittorCV/jittordet
import jittor as jt
from jittor import nn
from jittor.misc import _pair

__all__ = ['ROIAlign']

CUDA_HEADER = r'''
#include <cmath>
#include <cstdio>
#include <climits>
#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)

#define THREADS_PER_BLOCK 1024

inline int GET_BLOCKS(const int N) {
    int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    int max_block_num = 65000;
    return min(optimal_block_num, max_block_num);
}

template <typename scalar_t>
__device__ scalar_t bilinear_interpolate(const scalar_t *bottom_data,
                                         const int height, const int width,
                                         scalar_t y, scalar_t x) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    return 0;
  }

  if (y < 0) y = 0;
  if (x < 0) x = 0;

  int y_low = (int)y;
  int x_low = (int)x;
  int y_high;
  int x_high;

  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (scalar_t)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (scalar_t)x_low;
  } else {
    x_high = x_low + 1;
  }

  scalar_t ly = y - y_low;
  scalar_t lx = x - x_low;
  scalar_t hy = 1. - ly;
  scalar_t hx = 1. - lx;
  // do bilinear interpolation
  scalar_t lt = bottom_data[y_low * width + x_low];
  scalar_t rt = bottom_data[y_low * width + x_high];
  scalar_t lb = bottom_data[y_high * width + x_low];
  scalar_t rb = bottom_data[y_high * width + x_high];
  scalar_t w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  scalar_t val = (w1 * lt + w2 * rt + w3 * lb + w4 * rb);

  return val;
}

template <typename scalar_t>
__global__ void ROIAlignRotatedForward(const int nthreads, const scalar_t *bottom_data,
                                const scalar_t *bottom_rois,
                                const scalar_t spatial_scale,
                                const int sample_num, const int channels,
                                const int height, const int width,
                                const int pooled_height, const int pooled_width,
                                scalar_t *top_data) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;

    const scalar_t* offset_bottom_rois = bottom_rois + n * 6;
    int roi_batch_ind = offset_bottom_rois[0];

    // Do not using rounding; this implementation detail is critical
    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale - (scalar_t)0.5;
    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale - (scalar_t)0.5;
    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
    scalar_t theta = offset_bottom_rois[5];

    // Force malformed ROIs to be 1x1
    roi_width = max(roi_width, (scalar_t)1.);
    roi_height = max(roi_height, (scalar_t)1.);
    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) / static_cast<scalar_t>(pooled_height);
    scalar_t bin_size_w = static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);

    const scalar_t* offset_bottom_data =
        bottom_data + (roi_batch_ind * channels + c) * height * width;

    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sample_num > 0)
        ? sample_num
        : ceil(roi_height / pooled_height);  // e.g., = 2
    int roi_bin_grid_w =
        (sample_num > 0) ? sample_num : ceil(roi_width / pooled_width);

    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
    // Appropriate translation needs to be applied after.
    scalar_t roi_start_h = -roi_height / 2.0;
    scalar_t roi_start_w = -roi_width / 2.0;
    scalar_t cosscalar_theta = cos(theta);
    scalar_t sinscalar_theta = sin(theta);

    // We do average (integral) pooling inside a bin
    const scalar_t count = max(roi_bin_grid_h * roi_bin_grid_w, 1);  // e.g. = 4

    scalar_t output_val = 0.;
    for (int iy = 0; iy < roi_bin_grid_h; iy++) {  // e.g., iy = 0, 1
        const scalar_t yy = roi_start_h + ph * bin_size_h +
            static_cast<scalar_t>(iy + .5f) * bin_size_h /
                static_cast<scalar_t>(roi_bin_grid_h);  // e.g., 0.5, 1.5
        for (int ix = 0; ix < roi_bin_grid_w; ix++) {
        const scalar_t xx = roi_start_w + pw * bin_size_w +
            static_cast<scalar_t>(ix + .5f) * bin_size_w /
                static_cast<scalar_t>(roi_bin_grid_w);

        // Rotate by theta around the center and translate
        scalar_t x = xx * cosscalar_theta + yy * sinscalar_theta + roi_center_w;
        scalar_t y = yy * cosscalar_theta - xx * sinscalar_theta + roi_center_h;
        // scalar_t x = xx * cosscalar_theta - yy * sinscalar_theta + roi_center_w;
        // scalar_t y = xx * sinscalar_theta + yy * cosscalar_theta + roi_center_h;

        scalar_t val = bilinear_interpolate<scalar_t>(
            offset_bottom_data, height, width, y, x);
        output_val += val;
        }
    }
    output_val /= count;

    top_data[index] = output_val;
    }
}
template <typename scalar_t>
__device__ void bilinear_interpolate_gradient(const int height, const int width,
                                              scalar_t y, scalar_t x,
                                              scalar_t &w1, scalar_t &w2,
                                              scalar_t &w3, scalar_t &w4,
                                              int &x_low, int &x_high,
                                              int &y_low, int &y_high) {
  // deal with cases that inverse elements are out of feature map boundary
  if (y < -1.0 || y > height || x < -1.0 || x > width) {
    w1 = w2 = w3 = w4 = 0.;
    x_low = x_high = y_low = y_high = -1;
    return;
  }

  if (y < 0) y = 0;
  if (x < 0) x = 0;

  y_low = (int)y;
  x_low = (int)x;

  if (y_low >= height - 1) {
    y_high = y_low = height - 1;
    y = (scalar_t)y_low;
  } else {
    y_high = y_low + 1;
  }

  if (x_low >= width - 1) {
    x_high = x_low = width - 1;
    x = (scalar_t)x_low;
  } else {
    x_high = x_low + 1;
  }

  scalar_t ly = y - y_low;
  scalar_t lx = x - x_low;
  scalar_t hy = 1. - ly;
  scalar_t hx = 1. - lx;

  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

  return;
}

template <typename scalar_t>
__global__ void ROIAlignBackward(
    const int nthreads, const scalar_t *top_diff, const scalar_t *bottom_rois,
    const scalar_t spatial_scale, const int sample_num, const int channels,
    const int height, const int width, const int pooled_height,
    const int pooled_width, scalar_t *bottom_diff) {

    CUDA_1D_KERNEL_LOOP(index, nthreads) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;

    const scalar_t* offset_bottom_rois = bottom_rois + n * 6;
    int roi_batch_ind = offset_bottom_rois[0];

    // Do not round
    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale - (scalar_t)0.5;
    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale - (scalar_t)0.5;
    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
    scalar_t theta = offset_bottom_rois[5];


    // Force malformed ROIs to be 1x1
    roi_width = max(roi_width, (scalar_t)1.);
    roi_height = max(roi_height, (scalar_t)1.);
    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) / static_cast<scalar_t>(pooled_height);
    scalar_t bin_size_w = static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);

    scalar_t* offset_bottom_diff =
        bottom_diff + (roi_batch_ind * channels + c) * height * width;

    int top_offset = (n * channels + c) * pooled_height * pooled_width;
    const scalar_t* offset_top_diff = top_diff + top_offset;
    const scalar_t top_diff_this_bin = offset_top_diff[ph * pooled_width + pw];

    // We use roi_bin_grid to sample the grid and mimic integral
    int roi_bin_grid_h = (sample_num > 0)
        ? sample_num
        : ceil(roi_height / pooled_height);  // e.g., = 2
    int roi_bin_grid_w =
        (sample_num > 0) ? sample_num : ceil(roi_width / pooled_width);

    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
    // Appropriate translation needs to be applied after.
    scalar_t roi_start_h = -roi_height / 2.0;
    scalar_t roi_start_w = -roi_width / 2.0;
    scalar_t cosTheta = cos(theta);
    scalar_t sinTheta = sin(theta);

    // We do average (integral) pooling inside a bin
    const scalar_t count = roi_bin_grid_h * roi_bin_grid_w;  // e.g. = 4

    for (int iy = 0; iy < roi_bin_grid_h; iy++) {  // e.g., iy = 0, 1
      const scalar_t yy = roi_start_h + ph * bin_size_h +
          static_cast<scalar_t>(iy + .5f) * bin_size_h /
              static_cast<scalar_t>(roi_bin_grid_h);  // e.g., 0.5, 1.5
      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
        const scalar_t xx = roi_start_w + pw * bin_size_w +
            static_cast<scalar_t>(ix + .5f) * bin_size_w /
                static_cast<scalar_t>(roi_bin_grid_w);

        // Rotate by theta around the center and translate
        scalar_t x = xx * cosTheta + yy * sinTheta + roi_center_w;
        scalar_t y = yy * cosTheta - xx * sinTheta + roi_center_h;
        // scalar_t x = xx * cosTheta - yy * sinTheta + roi_center_w;
        // scalar_t y = xx * sinTheta + yy * cosTheta + roi_center_h;

        scalar_t w1, w2, w3, w4;
        int x_low, x_high, y_low, y_high;

        bilinear_interpolate_gradient<scalar_t>(
            height,
            width,
            y,
            x,
            w1,
            w2,
            w3,
            w4,
            x_low,
            x_high,
            y_low,
            y_high);

        scalar_t g1 = top_diff_this_bin * w1 / count;
        scalar_t g2 = top_diff_this_bin * w2 / count;
        scalar_t g3 = top_diff_this_bin * w3 / count;
        scalar_t g4 = top_diff_this_bin * w4 / count;

        if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
          atomicAdd(
              offset_bottom_diff + y_low * width + x_low, g1);
          atomicAdd(
              offset_bottom_diff + y_low * width + x_high, g2);
          atomicAdd(
              offset_bottom_diff + y_high * width + x_low, g3);
          atomicAdd(
              offset_bottom_diff + y_high * width + x_high, g4);
        }  // if
      }  // ix
    }  // iy
  }  // CUDA_1D_KERNEL_LOOP
}  // RoIAlignBackward
'''


class _RotatedROIAlign_v1(jt.Function):

    def execute(self, input, rois, output_size, spatial_scale, sampling_ratio):
        self.input = input
        self.rois = rois
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio
        assert rois.shape[1] == 6
        output_shapes = (rois.shape[0], input.shape[1], output_size[0],
                         output_size[1])
        return jt.code(
            output_shapes,
            input.dtype, [input, rois],
            cuda_header=CUDA_HEADER,
            cuda_src=f'''
                            @alias(input,in0);
                            @alias(rois,in1);
                            @alias(output,out0);
                            const float spatial_scale = {spatial_scale};
                            const float  sampling_ratio = {sampling_ratio};
                            auto num_rois = rois_shape0;
                            auto channels = input_shape1;
                            auto height = input_shape2;
                            auto width = input_shape3;
                            auto pooled_height = output_shape2;
                            auto pooled_width = output_shape3;
                            auto output_size = num_rois * pooled_height * pooled_width * channels;
                            ROIAlignRotatedForward<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>> (
                                        output_size, input_p, rois_p, spatial_scale,
                                        sampling_ratio, channels, height, width, pooled_height,
                                        pooled_width, output_p);
                            ''')

    def grad(self, output_grad):
        input, rois = self.input, self.rois
        input_grad = jt.code(
            input.shape,
            input.dtype, [input, rois, output_grad],
            cuda_header=CUDA_HEADER,
            cuda_src=f'''
                            @alias(input,in0)
                            @alias(rois,in1)
                            @alias(grad,in2)
                            @alias(grad_input,out0)
                            const float spatial_scale = {self.spatial_scale};
                            const float  sampling_ratio = {self.sampling_ratio};
                            auto num_rois = rois_shape0;
                            auto channels = input_shape1;
                            auto height = input_shape2;
                            auto width = input_shape3;
                            auto pooled_height = grad_shape2;
                            auto pooled_width = grad_shape3;
                            auto output_size = num_rois * pooled_height * pooled_width * channels;
                            cudaMemsetAsync(grad_input_p,0,grad_input->size);
                             ROIAlignBackward<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>>(
                                    output_size, grad_p, rois_p, spatial_scale, sampling_ratio,
                                    channels, height, width, pooled_height, pooled_width,
                                    grad_input_p);
                            ''')
        return input_grad, None


roi_align = _RotatedROIAlign_v1.apply


class ROIAlignRotated_v1(nn.Module):

    def __init__(self, output_size, spatial_scale, sampling_ratio=0):
        super(ROIAlignRotated_v1, self).__init__()
        self.output_size = _pair(output_size)
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio

    def execute(self, input, rois):
        return roi_align(input, rois, self.output_size, self.spatial_scale,
                         self.sampling_ratio)

    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 += ')'
        return tmpstr


def test_roialign():
    jt.flags.use_cuda = 1
    roi_align = ROIAlignRotated_v1((7, 7), 1 / 16.)
    feature = jt.randn((2, 1024, 64, 64))
    roi = jt.array([[0, 20, 120, 80, 195.5, 0.3], [1, 23, 56, 200, 300.5,
                                                   0.2]])
    output = roi_align(feature, roi)
    print(output.shape, output.sum())
    print(jt.grad(output.exp(), feature))


if __name__ == '__main__':
    test_roialign()

jittordet/ops/roi_align_rotated_v1.py
JittorCV/jittordet
import jittor as jt
from jittor import init, nn
from jittor.misc import _pair

__all__ = ['ROIPool']

CUDA_HEADER = r'''
#include <cmath>
#include <cstdio>
#include <climits>
using namespace std;
'''

CUDA_SRC = r'''
__global__ static void RoIPoolForwardKernel(@ARGS_DEF){
    @PRECALC
@alias(input,in0);
@alias(rois,in1);
@alias(output,out0);
@alias(argmax,out1);
const int pooled_height = output_shape2;
const int pooled_width = output_shape3;
const int  num_rois = rois_shape0;
const int  channels = input_shape1;
const int height = input_shape2;
const int width = input_shape3;
const int  nthreads = num_rois * pooled_height * pooled_width * channels;
const float spatial_scale = @in2(0);
auto bottom_data =  input_p;
auto bottom_rois = rois_p;
auto top_data = output_p;
auto argmax_data = argmax_p;
 for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x) {
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    auto offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    int roi_start_w = round(offset_bottom_rois[1] * spatial_scale);
    int roi_start_h = round(offset_bottom_rois[2] * spatial_scale);
    int roi_end_w = round(offset_bottom_rois[3] * spatial_scale);
    int roi_end_h = round(offset_bottom_rois[4] * spatial_scale);
    // Force malformed ROIs to be 1x1
    int roi_width = max(roi_end_w - roi_start_w + 1, 1);
    int roi_height = max(roi_end_h - roi_start_h + 1, 1);
    float bin_size_h = 1.0*roi_height / pooled_height;
    float bin_size_w = 1.0*roi_width / pooled_width;

    int hstart = static_cast<int>(floor(ph* bin_size_h));
    int wstart = static_cast<int>(floor(pw* bin_size_w));
    int hend = static_cast<int>(ceil((ph + 1)* bin_size_h));
    int wend = static_cast<int>(ceil((pw + 1)* bin_size_w));
    // Add roi offsets and clip to input boundaries
    hstart = min(max(hstart + roi_start_h, 0), height);
    hend = min(max(hend + roi_start_h, 0), height);
    wstart = min(max(wstart + roi_start_w, 0), width);
    wend = min(max(wend + roi_start_w, 0), width);
    bool is_empty = (hend <= hstart) || (wend <= wstart);
    // Define an empty pooling region to be zero
    float maxval = is_empty ? 0 : -99999999;
    // If nothing is pooled, argmax = -1 causes nothing to be backprop'd
    int maxidx = -1;
    auto offset_bottom_data =
        bottom_data + (roi_batch_ind * channels + c) * height * width;
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        int bottom_index = h * width + w;
        if (offset_bottom_data[bottom_index] > maxval) {
          maxval = offset_bottom_data[bottom_index];
          maxidx = bottom_index;
        }
      }
    }
    top_data[index] = maxval;
    argmax_data[index] = maxidx;
  }

}
cudaMemsetAsync(argmax_p,0,argmax->size);
cudaMemsetAsync(output_p,0,output->size);
const int total_count = in1_shape0 * out0_shape2 * out0_shape3 * in0_shape1;
const int thread_per_block = 1024;
const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
RoIPoolForwardKernel<<<block_count, thread_per_block>>>(@ARGS);
'''

CUDA_GRAD_SRC = [
    r'''
__global__ void RoIPoolBackwardKernel(@ARGS_DEF){
    @PRECALC
    @alias(input,in0);
    @alias(rois,in1);
    @alias(argmax,pout0)
    @alias(grad_input,out0);
    @alias(grad,dout);
    const float spatial_scale = @in2(0);
    const int pooled_height = pout0_shape2;
    const int pooled_width = pout0_shape3;
    const int num_rois = rois_shape0;
    const int channels = input_shape1;
    const int height = input_shape2;
    const int width = input_shape3;
    const int nthreads = num_rois * pooled_height * pooled_width * channels;
    auto top_diff = grad_p;
    auto argmax_data = argmax_p;
    auto bottom_diff = grad_input_p;
    auto bottom_rois = rois_p;
  for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x){
    // (n, c, ph, pw) is an element in the pooled output
    int pw = index % pooled_width;
    int ph = (index / pooled_width) % pooled_height;
    int c = (index / pooled_width / pooled_height) % channels;
    int n = index / pooled_width / pooled_height / channels;
    auto offset_bottom_rois = bottom_rois + n * 5;
    int roi_batch_ind = offset_bottom_rois[0];
    int bottom_offset = (roi_batch_ind * channels + c) * height * width;
    int top_offset    = (n * channels + c) * pooled_height * pooled_width;
    auto offset_top_diff = top_diff + top_offset;
    auto offset_bottom_diff = bottom_diff + bottom_offset;
    auto offset_argmax_data = argmax_data + top_offset;
    int argmax = offset_argmax_data[ph * pooled_width + pw];
    if (argmax != -1) {
      atomicAdd(
          offset_bottom_diff + argmax,
          offset_top_diff[ph * pooled_width + pw]);
    }
  }
}
cudaMemsetAsync(out0_p,0,out0->size);
const int total_count = in1_shape0*in0_shape1*pout0_shape2*pout0_shape3;
const int thread_per_block = 1024;
const int block_count = (total_count + thread_per_block - 1) / thread_per_block;
RoIPoolBackwardKernel<<<block_count, thread_per_block>>>(@ARGS);
''', r'''
''', r'''
'''
]


def roi_pool(input, rois, output_size, spatial_scale):
    output_size = _pair(output_size)
    spatial_scale = jt.array([spatial_scale])
    output_shapes = [
        (rois.shape[0], input.shape[1], output_size[0], output_size[1])
    ] * 2
    inputs = [input, rois, spatial_scale]
    output_types = [input.dtype, 'int32']
    output, arg_output = jt.code(
        output_shapes,
        output_types,
        inputs,
        cuda_header=CUDA_HEADER,
        cuda_src=CUDA_SRC,
        cuda_grad_src=CUDA_GRAD_SRC)
    return output


class ROIPool(nn.Module):

    def __init__(self, output_size, spatial_scale):
        super(ROIPool, self).__init__()
        self.output_size = output_size
        self.spatial_scale = spatial_scale

    def execute(self, input, rois):
        return roi_pool(input, rois, self.output_size, self.spatial_scale)

    def __repr__(self):
        tmpstr = self.__class__.__name__ + '('
        tmpstr += 'output_size=' + str(self.output_size)
        tmpstr += ', spatial_scale=' + str(self.spatial_scale)
        tmpstr += ')'
        return tmpstr

jittordet/ops/roi_pool.py
JittorCV/jittordet
from .base_data_element import BaseDataElement
from .det_data_sample import DetDataSample, OptSampleList, SampleList
from .instance_data import InstanceData, InstanceList, OptInstanceList

__all__ = [
    'BaseDataElement', 'InstanceData', 'DetDataSampler', 'DetDataSample',
    'SampleList', 'OptSampleList', 'InstanceList', 'OptInstanceList'
]

jittordet/structures/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab mmengine.structures.base_data_element
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import sys
from typing import Any, Iterator, Optional, Tuple, Type, Union

import jittor as jt
import numpy as np


class BaseDataElement:

    def __init__(self, *, metainfo: Optional[dict] = None, **kwargs) -> None:

        self._metainfo_fields: set = set()
        self._data_fields: set = set()

        if metainfo is not None:
            self.set_metainfo(metainfo=metainfo)
        if kwargs:
            self.set_data(kwargs)

    def set_metainfo(self, metainfo: dict) -> None:
        """Set or change key-value pairs in ``metainfo_field`` by parameter
        ``metainfo``.

        Args:
            metainfo (dict): A dict contains the meta information
                of image, such as ``img_shape``, ``scale_factor``, etc.
        """
        assert isinstance(
            metainfo,
            dict), f'metainfo should be a ``dict`` but got {type(metainfo)}'
        meta = copy.deepcopy(metainfo)
        for k, v in meta.items():
            self.set_field(name=k, value=v, field_type='metainfo', dtype=None)

    def set_data(self, data: dict) -> None:
        """Set or change key-value pairs in ``data_field`` by parameter
        ``data``.

        Args:
            data (dict): A dict contains annotations of image or
                model predictions.
        """
        assert isinstance(data,
                          dict), f'data should be a `dict` but got {data}'
        for k, v in data.items():
            # Use `setattr()` rather than `self.set_field` to allow `set_data`
            # to set property method.
            setattr(self, k, v)

    def update(self, instance: 'BaseDataElement') -> None:
        """The update() method updates the BaseDataElement with the elements
        from another BaseDataElement object.

        Args:
            instance (BaseDataElement): Another BaseDataElement object for
                update the current object.
        """
        assert isinstance(
            instance, BaseDataElement
        ), f'instance should be a `BaseDataElement` but got {type(instance)}'
        self.set_metainfo(dict(instance.metainfo_items()))
        self.set_data(dict(instance.items()))

    def new(self,
            *,
            metainfo: Optional[dict] = None,
            **kwargs) -> 'BaseDataElement':
        """Return a new data element with same type.

        If ``metainfo`` and
        ``data`` are None, the new data element will have same metainfo and
        data. If metainfo or data is not None, the new result will overwrite it
        with the input value.
        Args:
            metainfo (dict, optional): A dict contains the meta information
                of image, such as ``img_shape``, ``scale_factor``, etc.
                Defaults to None.
            kwargs (dict): A dict contains annotations of image or
                model predictions.
        Returns:
            BaseDataElement: A new data element with same type.
        """
        new_data = self.__class__()

        if metainfo is not None:
            new_data.set_metainfo(metainfo)
        else:
            new_data.set_metainfo(dict(self.metainfo_items()))
        if kwargs:
            new_data.set_data(kwargs)
        else:
            new_data.set_data(dict(self.items()))
        return new_data

    def clone(self):
        """Deep copy the current data element.

        Returns:
            BaseDataElement: The copy of current data element.
        """
        clone_data = self.__class__()
        clone_data.set_metainfo(dict(self.metainfo_items()))
        clone_data.set_data(dict(self.items()))
        return clone_data

    def keys(self) -> list:
        """
        Returns:
            list: Contains all keys in data_fields.
        """
        return list(self._data_fields)

    def metainfo_keys(self) -> list:
        """
        Returns:
            list: Contains all keys in metainfo_fields.
        """
        return list(self._metainfo_fields)

    def values(self) -> list:
        """
        Returns:
            list: Contains all values in data.
        """
        return [getattr(self, k) for k in self.keys()]

    def metainfo_values(self) -> list:
        """
        Returns:
            list: Contains all values in metainfo.
        """
        return [getattr(self, k) for k in self.metainfo_keys()]

    def all_keys(self) -> list:
        """
        Returns:
            list: Contains all keys in metainfo and data.
        """
        return self.metainfo_keys() + self.keys()

    def all_values(self) -> list:
        """
        Returns:
            list: Contains all values in metainfo and data.
        """
        return self.metainfo_values() + self.values()

    def all_items(self) -> Iterator[Tuple[str, Any]]:
        """
        Returns:
            iterator: An iterator object whose element is (key, value) tuple
            pairs for ``metainfo`` and ``data``.
        """
        for k in self.all_keys():
            yield (k, getattr(self, k))

    def items(self) -> Iterator[Tuple[str, Any]]:
        """
        Returns:
            iterator: An iterator object whose element is (key, value) tuple
            pairs for ``data``.
        """
        for k in self.keys():
            yield (k, getattr(self, k))

    def metainfo_items(self) -> Iterator[Tuple[str, Any]]:
        """
        Returns:
            iterator: An iterator object whose element is (key, value) tuple
            pairs for ``metainfo``.
        """
        for k in self.metainfo_keys():
            yield (k, getattr(self, k))

    @property
    def metainfo(self) -> dict:
        """dict: A dict contains metainfo of current data element."""
        return dict(self.metainfo_items())

    def __setattr__(self, name: str, value: Any):
        """setattr is only used to set data."""
        if name in ('_metainfo_fields', '_data_fields'):
            if not hasattr(self, name):
                super().__setattr__(name, value)
            else:
                raise AttributeError(f'{name} has been used as a '
                                     'private attribute, which is immutable.')
        else:
            self.set_field(
                name=name, value=value, field_type='data', dtype=None)

    def __delattr__(self, item: str):
        """Delete the item in dataelement.

        Args:
            item (str): The key to delete.
        """
        if item in ('_metainfo_fields', '_data_fields'):
            raise AttributeError(f'{item} has been used as a '
                                 'private attribute, which is immutable.')
        super().__delattr__(item)
        if item in self._metainfo_fields:
            self._metainfo_fields.remove(item)
        elif item in self._data_fields:
            self._data_fields.remove(item)

    # dict-like methods
    __delitem__ = __delattr__

    def get(self, key, default=None) -> Any:
        """Get property in data and metainfo as the same as python."""
        # Use `getattr()` rather than `self.__dict__.get()` to allow getting
        # properties.
        return getattr(self, key, default)

    def pop(self, *args) -> Any:
        """Pop property in data and metainfo as the same as python."""
        assert len(args) < 3, '``pop`` get more than 2 arguments'
        name = args[0]
        if name in self._metainfo_fields:
            self._metainfo_fields.remove(args[0])
            return self.__dict__.pop(*args)

        elif name in self._data_fields:
            self._data_fields.remove(args[0])
            return self.__dict__.pop(*args)

        # with default value
        elif len(args) == 2:
            return args[1]
        else:
            # don't just use 'self.__dict__.pop(*args)' for only popping key in
            # metainfo or data
            raise KeyError(f'{args[0]} is not contained in metainfo or data')

    def __contains__(self, item: str) -> bool:
        """Whether the item is in dataelement.

        Args:
            item (str): The key to inquire.
        """
        return item in self._data_fields or item in self._metainfo_fields

    def set_field(self,
                  value: Any,
                  name: str,
                  dtype: Optional[Union[Type, Tuple[Type, ...]]] = None,
                  field_type: str = 'data') -> None:
        """Special method for set union field, used as property.setter
        functions."""
        assert field_type in ['metainfo', 'data']
        if dtype is not None:
            assert isinstance(
                value,
                dtype), f'{value} should be a {dtype} but got {type(value)}'

        if field_type == 'metainfo':
            if name in self._data_fields:
                raise AttributeError(
                    f'Cannot set {name} to be a field of metainfo '
                    f'because {name} is already a data field')
            self._metainfo_fields.add(name)
        else:
            if name in self._metainfo_fields:
                raise AttributeError(
                    f'Cannot set {name} to be a field of data '
                    f'because {name} is already a metainfo field')
            # The name only added to `data_fields` when it is not the
            # attribute related to property(methods decorated by @property).
            if not isinstance(
                    getattr(type(self),
                            sys._getframe(1).f_code.co_name, None), property):
                self._data_fields.add(name)
        super().__setattr__(name, value)

    # Tensor-like methods
    def detach(self) -> 'BaseDataElement':
        """Detach all tensors in data."""
        new_data = self.new()
        for k, v in self.items():
            if isinstance(v, (jt.Var, BaseDataElement)):
                v = v.detach()
                data = {k: v}
                new_data.set_data(data)
        return new_data

    # Tensor-like methods
    def numpy(self) -> 'BaseDataElement':
        """Convert all tensors to np.ndarray in data."""
        new_data = self.new()
        for k, v in self.items():
            if isinstance(v, (jt.Var, BaseDataElement)):
                v = v.detach().numpy()
                data = {k: v}
                new_data.set_data(data)
        return new_data

    def to_jittor(self, stop_grad=True) -> 'BaseDataElement':
        """Convert all np.ndarray to tensor in data."""
        new_data = self.new()
        for k, v in self.items():
            data = {}
            if isinstance(v, np.ndarray):
                v = jt.array(v)
                if stop_grad:
                    v = v.stop_grad()
                data[k] = v
            elif isinstance(v, BaseDataElement):
                v = v.to_jittor(stop_grad)
                data[k] = v
            new_data.set_data(data)
        return new_data

    def to_dict(self) -> dict:
        """Convert BaseDataElement to dict."""
        return {
            k: v.to_dict() if isinstance(v, BaseDataElement) else v
            for k, v in self.all_items()
        }

    def __repr__(self) -> str:
        """Represent the object."""

        def _addindent(s_: str, num_spaces: int) -> str:
            """This func is modified from `pytorch` https://github.com/pytorch/
            pytorch/blob/b17b2b1cc7b017c3daaeff8cc7ec0f514d42ec37/torch/nn/modu
            les/module.py#L29.

            Args:
                s_ (str): The string to add spaces.
                num_spaces (int): The num of space to add.
            Returns:
                str: The string after add indent.
            """
            s = s_.split('\n')
            # don't do anything for single-line stuff
            if len(s) == 1:
                return s_
            first = s.pop(0)
            s = [(num_spaces * ' ') + line for line in s]
            s = '\n'.join(s)  # type: ignore
            s = first + '\n' + s  # type: ignore
            return s  # type: ignore

        def dump(obj: Any) -> str:
            """Represent the object.

            Args:
                obj (Any): The obj to represent.
            Returns:
                str: The represented str.
            """
            _repr = ''
            if isinstance(obj, dict):
                for k, v in obj.items():
                    _repr += f'\n{k}: {_addindent(dump(v), 4)}'
            elif isinstance(obj, BaseDataElement):
                _repr += '\n\n    META INFORMATION'
                metainfo_items = dict(obj.metainfo_items())
                _repr += _addindent(dump(metainfo_items), 4)
                _repr += '\n\n    DATA FIELDS'
                items = dict(obj.items())
                _repr += _addindent(dump(items), 4)
                classname = obj.__class__.__name__
                _repr = f'<{classname}({_repr}\n) at {hex(id(obj))}>'
            else:
                _repr += repr(obj)
            return _repr

        return dump(self)

jittordet/structures/base_data_element.py
JittorCV/jittordet
# Modified from OpenMMLab mmdet/structures/det_data_sample.py
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional

from .base_data_element import BaseDataElement
from .instance_data import InstanceData


class DetDataSample(BaseDataElement):

    @property
    def proposals(self) -> InstanceData:
        return self._proposals

    @proposals.setter
    def proposals(self, value: InstanceData):
        self.set_field(value, '_proposals', dtype=InstanceData)

    @proposals.deleter
    def proposals(self):
        del self._proposals

    @property
    def gt_instances(self) -> InstanceData:
        return self._gt_instances

    @gt_instances.setter
    def gt_instances(self, value: InstanceData):
        self.set_field(value, '_gt_instances', dtype=InstanceData)

    @gt_instances.deleter
    def gt_instances(self):
        del self._gt_instances

    @property
    def pred_instances(self) -> InstanceData:
        return self._pred_instances

    @pred_instances.setter
    def pred_instances(self, value: InstanceData):
        self.set_field(value, '_pred_instances', dtype=InstanceData)

    @pred_instances.deleter
    def pred_instances(self):
        del self._pred_instances

    @property
    def ignored_instances(self) -> InstanceData:
        return self._ignored_instances

    @ignored_instances.setter
    def ignored_instances(self, value: InstanceData):
        self.set_field(value, '_ignored_instances', dtype=InstanceData)

    @ignored_instances.deleter
    def ignored_instances(self):
        del self._ignored_instances


SampleList = List[DetDataSample]
OptSampleList = Optional[SampleList]

jittordet/structures/det_data_sample.py
JittorCV/jittordet
# Modified from OpenMMLab mmengine.structures.instance_data.py
# Copyright (c) OpenMMLab. All rights reserved.
import itertools
from collections.abc import Sized
from typing import List, Optional

import jittor as jt
import numpy as np

from .base_data_element import BaseDataElement


class InstanceData(BaseDataElement):

    def __setattr__(self, name: str, value: Sized):
        """setattr is only used to set data.

        The value must have the attribute of `__len__` and have the same length
        of `InstanceData`.
        """
        if name in ('_metainfo_fields', '_data_fields'):
            if not hasattr(self, name):
                super().__setattr__(name, value)
            else:
                raise AttributeError(f'{name} has been used as a '
                                     'private attribute, which is immutable.')

        else:
            assert isinstance(value,
                              Sized), 'value must contain `__len__` attribute'

            if len(self) > 0:
                assert len(value) == len(self), 'The length of ' \
                                                f'values {len(value)} is ' \
                                                'not consistent with ' \
                                                'the length of this ' \
                                                ':obj:`InstanceData` ' \
                                                f'{len(self)}'
            super().__setattr__(name, value)

    __setitem__ = __setattr__

    def __getitem__(self, item) -> 'InstanceData':
        """
        Args:
            item (str, int, list, :obj:`slice`, :obj:`numpy.ndarray`,
                :obj:`torch.LongTensor`, :obj:`torch.BoolTensor`):
                Get the corresponding values according to item.
        Returns:
            :obj:`InstanceData`: Corresponding values.
        """
        if isinstance(item, list):
            item = np.array(item)
        if isinstance(item, np.ndarray):
            # The default int type of numpy is platform dependent, int32 for
            # windows and int64 for linux. `torch.Tensor` requires the index
            # should be int64, therefore we simply convert it to int64 here.
            # More details in https://github.com/numpy/numpy/issues/9464
            item = item.astype(np.int64) if item.dtype == np.int32 else item
            item = jt.array(item)

        if isinstance(item, str):
            return getattr(self, item)

        if isinstance(item, int):
            if item >= len(self) or item < -len(self):  # type:ignore
                raise IndexError(f'Index {item} out of range!')
            else:
                # keep the dimension
                item = slice(item, None, len(self))

        new_data = self.__class__(metainfo=self.metainfo)
        if isinstance(item, jt.Var):
            assert item.ndim == 1, 'Only support to get the' \
                                   ' values along the first dimension.'

            for k, v in self.items():
                if isinstance(v, jt.Var):
                    new_data[k] = v[item]
                elif isinstance(v, np.ndarray):
                    new_data[k] = v[item.numpy()]
                elif isinstance(
                        v, (str, list, tuple)) or (hasattr(v, '__getitem__')
                                                   and hasattr(v, 'cat')):
                    # convert to indexes from BoolTensor
                    if item.dtype == jt.bool:
                        indexes = jt.nonzero(item).view(-1).numpy().tolist()
                    else:
                        indexes = item.numpy().tolist()
                    slice_list = []
                    if indexes:
                        for index in indexes:
                            slice_list.append(slice(index, None, len(v)))
                    else:
                        slice_list.append(slice(None, 0, None))
                    r_list = [v[s] for s in slice_list]
                    if isinstance(v, (str, list, tuple)):
                        new_value = r_list[0]
                        for r in r_list[1:]:
                            new_value = new_value + r
                    else:
                        new_value = v.cat(r_list)
                    new_data[k] = new_value
                else:
                    raise ValueError(
                        f'The type of `{k}` is `{type(v)}`, which has no '
                        'attribute of `cat`, so it does not '
                        'support slice with `bool`')

        else:
            # item is a slice
            for k, v in self.items():
                new_data[k] = v[item]
        return new_data  # type:ignore

    @staticmethod
    def cat(instances_list: List['InstanceData']) -> 'InstanceData':
        """Concat the instances of all :obj:`InstanceData` in the list.

        Note: To ensure that cat returns as expected, make sure that
        all elements in the list must have exactly the same keys.
        Args:
            instances_list (list[:obj:`InstanceData`]): A list
                of :obj:`InstanceData`.
        Returns:
            :obj:`InstanceData`
        """
        assert all(
            isinstance(results, InstanceData) for results in instances_list)
        assert len(instances_list) > 0
        if len(instances_list) == 1:
            return instances_list[0]

        # metainfo and data_fields must be exactly the
        # same for each element to avoid exceptions.
        field_keys_list = [
            instances.all_keys() for instances in instances_list
        ]
        assert len({len(field_keys) for field_keys in field_keys_list}) \
               == 1 and len(set(itertools.chain(*field_keys_list))) \
               == len(field_keys_list[0]), 'There are different keys in ' \
                                           '`instances_list`, which may ' \
                                           'cause the cat operation ' \
                                           'to fail. Please make sure all ' \
                                           'elements in `instances_list` ' \
                                           'have the exact same key.'

        new_data = instances_list[0].__class__(
            metainfo=instances_list[0].metainfo)
        for k in instances_list[0].keys():
            values = [results[k] for results in instances_list]
            v0 = values[0]
            if isinstance(v0, jt.Var):
                new_values = jt.cat(values, dim=0)
            elif isinstance(v0, np.ndarray):
                new_values = np.concatenate(values, axis=0)
            elif isinstance(v0, (str, list, tuple)):
                new_values = v0[:]
                for v in values[1:]:
                    new_values += v
            elif hasattr(v0, 'cat'):
                new_values = v0.cat(values)
            else:
                raise ValueError(
                    f'The type of `{k}` is `{type(v0)}` which has no '
                    'attribute of `cat`')
            new_data[k] = new_values
        return new_data  # type:ignore

    def __len__(self) -> int:
        """int: The length of InstanceData."""
        if len(self._data_fields) > 0:
            return len(self.values()[0])
        else:
            return 0


InstanceList = List[InstanceData]
OptInstanceList = Optional[InstanceList]

jittordet/structures/instance_data.py
JittorCV/jittordet
from .bbox_overlaps import bbox_overlaps
from .image import (_scale_size, imflip, imnormalize, impad, impad_to_multiple,
                    imrescale, imresize, rescale_size)
from .types import is_list_of, is_seq_of, is_tuple_of
from .util_random import ensure_rng

__all__ = [
    '_scale_size', 'imresize', 'rescale_size', 'imrescale', 'imflip',
    'imnormalize', 'impad', 'impad_to_multiple', 'is_seq_of', 'is_list_of',
    'is_tuple_of', 'bbox_overlaps', 'ensure_rng'
]

jittordet/utils/__init__.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/structures/bbox/bbox_overlaps.py
# Copyright (c) OpenMMLab. All rights reserved.
import jittor as jt


def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
    """Calculate overlap between two set of bboxes.

    FP16 Contributed by https://github.com/open-mmlab/mmdetection/pull/4889
    Note:
        Assume bboxes1 is M x 4, bboxes2 is N x 4, when mode is 'iou',
        there are some new generated variable when calculating IOU
        using bbox_overlaps function:

        1) is_aligned is False
            area1: M x 1
            area2: N x 1
            lt: M x N x 2
            rb: M x N x 2
            wh: M x N x 2
            overlap: M x N x 1
            union: M x N x 1
            ious: M x N x 1

            Total memory:
                S = (9 x N x M + N + M) * 4 Byte,

            When using FP16, we can reduce:
                R = (9 x N x M + N + M) * 4 / 2 Byte
                R large than (N + M) * 4 * 2 is always true when N and M >= 1.
                Obviously, N + M <= N * M < 3 * N * M, when N >=2 and M >=2,
                           N + 1 < 3 * N, when N or M is 1.

            Given M = 40 (ground truth), N = 400000 (three anchor boxes
            in per grid, FPN, R-CNNs),
                R = 275 MB (one times)

            A special case (dense detection), M = 512 (ground truth),
                R = 3516 MB = 3.43 GB

            When the batch size is B, reduce:
                B x R

            Therefore, CUDA memory runs out frequently.

            Experiments on GeForce RTX 2080Ti (11019 MiB):

            |   dtype   |   M   |   N   |   Use    |   Real   |   Ideal   |
            |:----:|:----:|:----:|:----:|:----:|:----:|
            |   FP32   |   512 | 400000 | 8020 MiB |   --   |   --   |
            |   FP16   |   512 | 400000 |   4504 MiB | 3516 MiB | 3516 MiB |
            |   FP32   |   40 | 400000 |   1540 MiB |   --   |   --   |
            |   FP16   |   40 | 400000 |   1264 MiB |   276MiB   | 275 MiB |

        2) is_aligned is True
            area1: N x 1
            area2: N x 1
            lt: N x 2
            rb: N x 2
            wh: N x 2
            overlap: N x 1
            union: N x 1
            ious: N x 1

            Total memory:
                S = 11 x N * 4 Byte

            When using FP16, we can reduce:
                R = 11 x N * 4 / 2 Byte

        So do the 'giou' (large than 'iou').

        Time-wise, FP16 is generally faster than FP32.

        When gpu_assign_thr is not -1, it takes more time on cpu
        but not reduce memory.
        There, we can reduce half the memory and keep the speed.

    If ``is_aligned`` is ``False``, then calculate the overlaps between each
    bbox of bboxes1 and bboxes2, otherwise the overlaps between each aligned
    pair of bboxes1 and bboxes2.

    Args:
        bboxes1 (Tensor): shape (B, m, 4) in <x1, y1, x2, y2> format or empty.
        bboxes2 (Tensor): shape (B, n, 4) in <x1, y1, x2, y2> format or empty.
            B indicates the batch dim, in shape (B1, B2, ..., Bn).
            If ``is_aligned`` is ``True``, then m and n must be equal.
        mode (str): "iou" (intersection over union), "iof" (intersection over
            foreground) or "giou" (generalized intersection over union).
            Default "iou".
        is_aligned (bool, optional): If True, then m and n must be equal.
            Default False.
        eps (float, optional): A value added to the denominator for numerical
            stability. Default 1e-6.

    Returns:
        Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)

    Example:
        >>> bboxes1 = torch.FloatTensor([
        >>>     [0, 0, 10, 10],
        >>>     [10, 10, 20, 20],
        >>>     [32, 32, 38, 42],
        >>> ])
        >>> bboxes2 = torch.FloatTensor([
        >>>     [0, 0, 10, 20],
        >>>     [0, 10, 10, 19],
        >>>     [10, 10, 20, 20],
        >>> ])
        >>> overlaps = bbox_overlaps(bboxes1, bboxes2)
        >>> assert overlaps.shape == (3, 3)
        >>> overlaps = bbox_overlaps(bboxes1, bboxes2, is_aligned=True)
        >>> assert overlaps.shape == (3, )

    Example:
        >>> empty = torch.empty(0, 4)
        >>> nonempty = torch.FloatTensor([[0, 0, 10, 9]])
        >>> assert tuple(bbox_overlaps(empty, nonempty).shape) == (0, 1)
        >>> assert tuple(bbox_overlaps(nonempty, empty).shape) == (1, 0)
        >>> assert tuple(bbox_overlaps(empty, empty).shape) == (0, 0)
    """

    assert mode in ['iou', 'iof', 'giou'], f'Unsupported mode {mode}'
    # Either the boxes are empty or the length of boxes' last dimension is 4
    assert (bboxes1.size(-1) == 4 or bboxes1.size(0) == 0)
    assert (bboxes2.size(-1) == 4 or bboxes2.size(0) == 0)

    # Batch dim must be the same
    # Batch dim: (B1, B2, ... Bn)
    assert bboxes1.shape[:-2] == bboxes2.shape[:-2]
    batch_shape = bboxes1.shape[:-2]

    rows = bboxes1.size(-2)
    cols = bboxes2.size(-2)
    if is_aligned:
        assert rows == cols

    if rows * cols == 0:
        if is_aligned:
            return bboxes1.new(batch_shape + (rows, ))
        else:
            return bboxes1.new(batch_shape + (rows, cols))

    area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
        bboxes1[..., 3] - bboxes1[..., 1])
    area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
        bboxes2[..., 3] - bboxes2[..., 1])

    if is_aligned:
        lt = jt.maximum(bboxes1[..., :2], bboxes2[..., :2])  # [B, rows, 2]
        rb = jt.minimum(bboxes1[..., 2:], bboxes2[..., 2:])  # [B, rows, 2]

        wh = (rb - lt).clamp(min_v=0)
        overlap = wh[..., 0] * wh[..., 1]

        if mode in ['iou', 'giou']:
            union = area1 + area2 - overlap
        else:
            union = area1
        if mode == 'giou':
            enclosed_lt = jt.minimum(bboxes1[..., :2], bboxes2[..., :2])
            enclosed_rb = jt.maximum(bboxes1[..., 2:], bboxes2[..., 2:])
    else:
        lt = jt.maximum(bboxes1[..., :, None, :2],
                        bboxes2[..., None, :, :2])  # [B, rows, cols, 2]
        rb = jt.minimum(bboxes1[..., :, None, 2:],
                        bboxes2[..., None, :, 2:])  # [B, rows, cols, 2]

        wh = (rb - lt).clamp(min_v=0)
        overlap = wh[..., 0] * wh[..., 1]

        if mode in ['iou', 'giou']:
            union = area1[..., None] + area2[..., None, :] - overlap
        else:
            union = area1[..., None]
        if mode == 'giou':
            enclosed_lt = jt.minimum(bboxes1[..., :, None, :2],
                                     bboxes2[..., None, :, :2])
            enclosed_rb = jt.maximum(bboxes1[..., :, None, 2:],
                                     bboxes2[..., None, :, 2:])

    union = union.clamp(min_v=eps)
    ious = overlap / union
    if mode in ['iou', 'iof']:
        return ious
    # calculate gious
    enclose_wh = (enclosed_rb - enclosed_lt).clamp(min_v=0)
    enclose_area = enclose_wh[..., 0] * enclose_wh[..., 1]
    enclose_area = enclose_area.clamp(min_v=eps)
    gious = ious - (enclose_area - union) / enclose_area
    return gious

jittordet/utils/bbox_overlaps.py
JittorCV/jittordet
import numbers
from typing import List, Optional, Tuple, Union

import cv2
import numpy as np

try:
    from PIL import Image
except ImportError:
    Image = None

imread_backend = 'cv2'


def _scale_size(
    size: Tuple[int, int],
    scale: Union[float, int, tuple],
) -> Tuple[int, int]:
    """Rescale a size by a ratio."""
    if isinstance(scale, (float, int)):
        scale = (scale, scale)
    w, h = size
    return int(w * float(scale[0]) + 0.5), int(h * float(scale[1]) + 0.5)


cv2_interp_codes = {
    'nearest': cv2.INTER_NEAREST,
    'bilinear': cv2.INTER_LINEAR,
    'bicubic': cv2.INTER_CUBIC,
    'area': cv2.INTER_AREA,
    'lanczos': cv2.INTER_LANCZOS4
}

cv2_border_modes = {
    'constant': cv2.BORDER_CONSTANT,
    'replicate': cv2.BORDER_REPLICATE,
    'reflect': cv2.BORDER_REFLECT,
    'wrap': cv2.BORDER_WRAP,
    'reflect_101': cv2.BORDER_REFLECT_101,
    'transparent': cv2.BORDER_TRANSPARENT,
    'isolated': cv2.BORDER_ISOLATED
}

# Pillow >=v9.1.0 use a slightly different naming scheme for filters.
# Set pillow_interp_codes according to the naming scheme used.
if Image is not None:
    if hasattr(Image, 'Resampling'):
        pillow_interp_codes = {
            'nearest': Image.Resampling.NEAREST,
            'bilinear': Image.Resampling.BILINEAR,
            'bicubic': Image.Resampling.BICUBIC,
            'box': Image.Resampling.BOX,
            'lanczos': Image.Resampling.LANCZOS,
            'hamming': Image.Resampling.HAMMING
        }
    else:
        pillow_interp_codes = {
            'nearest': Image.NEAREST,
            'bilinear': Image.BILINEAR,
            'bicubic': Image.BICUBIC,
            'box': Image.BOX,
            'lanczos': Image.LANCZOS,
            'hamming': Image.HAMMING
        }


def imresize(
    img: np.ndarray,
    size: Tuple[int, int],
    return_scale: bool = False,
    interpolation: str = 'bilinear',
    out: Optional[np.ndarray] = None,
    backend: Optional[str] = None
) -> Union[Tuple[np.ndarray, float, float], np.ndarray]:
    """Resize image to a given size."""
    h, w = img.shape[:2]
    if backend is None:
        backend = imread_backend
    if backend not in ['cv2', 'pillow']:
        raise ValueError(f'backend: {backend} is not supported for resize.'
                         f"Supported backends are 'cv2', 'pillow'")

    if backend == 'pillow':
        assert img.dtype == np.uint8, 'Pillow backend only support uint8 type'
        pil_image = Image.fromarray(img)
        pil_image = pil_image.resize(size, pillow_interp_codes[interpolation])
        resized_img = np.array(pil_image)
    else:
        resized_img = cv2.resize(
            img, size, dst=out, interpolation=cv2_interp_codes[interpolation])
    if not return_scale:
        return resized_img
    else:
        w_scale = size[0] / w
        h_scale = size[1] / h
        return resized_img, w_scale, h_scale


def rescale_size(old_size: tuple,
                 scale: Union[float, int, tuple],
                 return_scale: bool = False) -> tuple:
    """Calculate the new size to be rescaled to."""
    w, h = old_size
    if isinstance(scale, (float, int)):
        if scale <= 0:
            raise ValueError(f'Invalid scale {scale}, must be positive.')
        scale_factor = scale
    elif isinstance(scale, tuple):
        max_long_edge = max(scale)
        max_short_edge = min(scale)
        scale_factor = min(max_long_edge / max(h, w),
                           max_short_edge / min(h, w))
    else:
        raise TypeError(
            f'Scale must be a number or tuple of int, but got {type(scale)}')

    new_size = _scale_size((w, h), scale_factor)

    if return_scale:
        return new_size, scale_factor
    else:
        return new_size


def imrescale(
    img: np.ndarray,
    scale: Union[float, Tuple[int, int]],
    return_scale: bool = False,
    interpolation: str = 'bilinear',
    backend: Optional[str] = None
) -> Union[np.ndarray, Tuple[np.ndarray, float]]:
    """Resize image while keeping the aspect ratio."""
    h, w = img.shape[:2]
    new_size, scale_factor = rescale_size((w, h), scale, return_scale=True)
    rescaled_img = imresize(
        img, new_size, interpolation=interpolation, backend=backend)
    if return_scale:
        return rescaled_img, scale_factor
    else:
        return rescaled_img


def imflip(img: np.ndarray, direction: str = 'horizontal') -> np.ndarray:
    """Flip an image horizontally or vertically."""
    assert direction in ['horizontal', 'vertical', 'diagonal']
    if direction == 'horizontal':
        return np.flip(img, axis=1)
    elif direction == 'vertical':
        return np.flip(img, axis=0)
    else:
        return np.flip(img, axis=(0, 1))


def imnormalize(img, mean, std, to_rgb=True):
    """Normalize an image with mean and std."""
    img = img.copy().astype(np.float32)
    assert img.dtype != np.uint8
    mean = np.float64(mean.reshape(1, -1))
    stdinv = 1 / np.float64(std.reshape(1, -1))
    if to_rgb:
        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
    cv2.subtract(img, mean, img)  # inplace
    cv2.multiply(img, stdinv, img)  # inplace
    return img


def impad(img: np.ndarray,
          *,
          shape: Optional[Tuple[int, int]] = None,
          padding: Union[int, tuple, None] = None,
          pad_val: Union[float, List] = 0,
          padding_mode: str = 'constant') -> np.ndarray:
    """Pad the given image to a certain shape or pad on all sides with
    specified padding mode and padding value."""

    assert (shape is not None) ^ (padding is not None)
    if shape is not None:
        width = max(shape[1] - img.shape[1], 0)
        height = max(shape[0] - img.shape[0], 0)
        padding = (0, 0, width, height)

    # check pad_val
    if isinstance(pad_val, tuple):
        assert len(pad_val) == img.shape[-1]
    elif not isinstance(pad_val, numbers.Number):
        raise TypeError('pad_val must be a int or a tuple. '
                        f'But received {type(pad_val)}')

    # check padding
    if isinstance(padding, tuple) and len(padding) in [2, 4]:
        if len(padding) == 2:
            padding = (padding[0], padding[1], padding[0], padding[1])
    elif isinstance(padding, numbers.Number):
        padding = (padding, padding, padding, padding)
    else:
        raise ValueError('Padding must be a int or a 2, or 4 element tuple.'
                         f'But received {padding}')

    # check padding mode
    assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']

    border_type = {
        'constant': cv2.BORDER_CONSTANT,
        'edge': cv2.BORDER_REPLICATE,
        'reflect': cv2.BORDER_REFLECT_101,
        'symmetric': cv2.BORDER_REFLECT
    }
    img = cv2.copyMakeBorder(
        img,
        padding[1],
        padding[3],
        padding[0],
        padding[2],
        border_type[padding_mode],
        value=pad_val)

    return img


def impad_to_multiple(img: np.ndarray,
                      divisor: int,
                      pad_val: Union[float, List] = 0) -> np.ndarray:
    """Pad an image to ensure each edge to be multiple to some number."""
    pad_h = int(np.ceil(img.shape[0] / divisor)) * divisor
    pad_w = int(np.ceil(img.shape[1] / divisor)) * divisor
    return impad(img, shape=(pad_h, pad_w), pad_val=pad_val)

jittordet/utils/image.py
JittorCV/jittordet
from collections import abc


def is_seq_of(seq, expected_type, seq_type=None):
    """Check whether it is a sequence of some type.

    Args:
        seq (Sequence): The sequence to be checked.
        expected_type (type or tuple): Expected type of sequence items.
        seq_type (type, optional): Expected sequence type. Defaults to None.

    Returns:
        bool: Return True if ``seq`` is valid else False.

    Examples:
        >>> from mmengine.utils import is_seq_of
        >>> seq = ['a', 'b', 'c']
        >>> is_seq_of(seq, str)
        True
        >>> is_seq_of(seq, int)
        False
    """
    if seq_type is None:
        exp_seq_type = abc.Sequence
    else:
        assert isinstance(seq_type, type)
        exp_seq_type = seq_type
    if not isinstance(seq, exp_seq_type):
        return False
    for item in seq:
        if not isinstance(item, expected_type):
            return False
    return True


def is_list_of(seq, expected_type):
    """Check whether it is a list of some type.

    A partial method of :func:`is_seq_of`.
    """
    return is_seq_of(seq, expected_type, seq_type=list)


def is_tuple_of(seq, expected_type):
    """Check whether it is a tuple of some type.

    A partial method of :func:`is_seq_of`.
    """
    return is_seq_of(seq, expected_type, seq_type=tuple)

jittordet/utils/types.py
JittorCV/jittordet
# Modified from OpenMMLab. mmdet/utils/util_random.py
# Copyright (c) OpenMMLab. All rights reserved.
"""Helpers for random number generators."""
import numpy as np


def ensure_rng(rng=None):
    """Coerces input into a random number generator.

    If the input is None, then a global random state is returned.

    If the input is a numeric value, then that is used as a seed to construct a
    random state. Otherwise the input is returned as-is.

    Adapted from [1]_.

    Args:
        rng (int | numpy.random.RandomState | None):
            if None, then defaults to the global rng. Otherwise this can be an
            integer or a RandomState class
    Returns:
        (numpy.random.RandomState) : rng -
            a numpy random number generator

    References:
        .. [1] https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270  # noqa: E501
    """

    if rng is None:
        rng = np.random.mtrand._rand
    elif isinstance(rng, int):
        rng = np.random.RandomState(rng)
    else:
        rng = rng
    return rng

jittordet/utils/util_random.py
JittorCV/jittordet
from setuptools import find_packages, setup


def readme():
    with open('README.md', encoding='utf-8') as f:
        content = f.read()
    return content


def get_version():
    with open('jittordet/version.py', 'r') as f:
        exec(compile(f.read(), 'jittordet/version.py', 'exec'))
    return locals()['__version__']


def get_requires():
    with open('requirements.txt', encoding='utf-8') as f:
        requires = [line for line in f]
    return requires


if __name__ == '__main__':
    setup(
        name='jittordet',
        version=get_version(),
        description='An object detection codebase based on Jittor.',
        long_description=readme(),
        long_description_content_type='text/markdown',
        packages=find_packages(exclude=('configs', 'tools')),
        classifiers=[
            'Development Status :: 3 - Alpha',
            'License :: OSI Approved :: Apache Software License',
            'Operating System :: OS Independent',
            'Programming Language :: Python :: 3',
            'Programming Language :: Python :: 3.8',
            'Programming Language :: Python :: 3.9',
        ],
        license='Apache License 2.0',
        install_requires=get_requires())

setup.py
JittorCV/jittordet
import argparse
import os.path as osp

from jittordet.engine import Runner, load_cfg


def parse_args():
    parser = argparse.ArgumentParser(description='Test a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('checkpoint', help='checkpoint to load from')
    parser.add_argument('--work-dir', help='the dir to save logs')
    parser.add_argument(
        '--disable-cuda',
        action='store_true',
        help='disable cuda and use cpu to train net.')
    return parser.parse_args()


def main():
    args = parse_args()
    cfg = load_cfg(args.config)

    cfg.load_from = args.checkpoint
    if args.work_dir is not None:
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    # set disable cuda
    cfg.disable_cuda = args.disable_cuda

    runner = Runner.from_cfg(cfg)
    runner.test()


if __name__ == '__main__':
    main()

tools/test.py
JittorCV/jittordet
import argparse
import os.path as osp

from jittordet.engine import Runner, load_cfg


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument('--seed', type=int, help='random seed for traning')
    parser.add_argument(
        '--disable-cuda',
        action='store_true',
        help='disable cuda and use cpu to train net.')
    return parser.parse_args()


def main():
    args = parse_args()
    cfg = load_cfg(args.config)

    if args.work_dir is not None:
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    # set seed
    if args.seed is not None:
        cfg.seed = args.seed
    # set disable cuda
    cfg.disable_cuda = args.disable_cuda

    runner = Runner.from_cfg(cfg)
    runner.train()


if __name__ == '__main__':
    main()

tools/train.py
JittorCV/jittordet