<reponame>LoSealL/centermask2
import cv2 as cv
import numpy as np
import torch
from detectron2.modeling import GeneralizedRCNN
from detectron2.layers import batched_nms
from torch.onnx import export

from train_net import DetectionCheckpointer, default_argument_parser, setup


class RCNNExporter(GeneralizedRCNN):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def preprocess_image(self, x):
        """
        Normalize, pad and batch the input images.
        """
        x -= self.pixel_mean.view(-1, 3, 1, 1)
        x /= self.pixel_std.view(-1, 3, 1, 1)
        return x

    def forward(self, inference_type, *x):
        if inference_type == 'fpn':
            inputs = self.preprocess_image(*x)
            features = self.backbone(inputs)
            return features
        elif inference_type == 'fcoshead':
            logits_pred, reg_pred, ctrness_pred, _ = self.proposal_generator.fcos_head(*x)
            strides = self.proposal_generator.strides
            bundle = (self.precal_locations, logits_pred, reg_pred, ctrness_pred, strides)
            scores = []
            centerness = []
            detections = []
            for l, o, r, c, s in zip(*bundle):
                r *= s
                N, C, H, W = o.shape
                box_cls = o.permute(0, 2, 3, 1).reshape(N, -1, C).sigmoid()
                box_regression = r.view(N, 4, H, W).permute(0, 2, 3, 1).reshape(N, -1, 4)
                ctrness = c.view(N, 1, H, W).permute(0, 2, 3, 1).reshape(N, -1).sigmoid()
                detection = torch.stack([
                    l[:, 0][None, ...] - box_regression[..., 0],
                    l[:, 1][None, ...] - box_regression[..., 1],
                    l[:, 0][None, ...] + box_regression[..., 2],
                    l[:, 1][None, ...] + box_regression[..., 3],
                ], dim=-1)
                scores.append(box_cls)
                centerness.append(ctrness)
                detections.append(detection)
            scores = torch.cat(scores, 1)
            centerness = torch.cat(centerness, 1)
            detections = torch.cat(detections, 1)
            return scores, detections / 128, centerness
        elif inference_type == 'centerhead':
            mask = self.roi_heads.mask_head(*x).sigmoid()
            return mask

    def export_fpn(self, input_shape, export_name):
        inputs = torch.randn(*input_shape)
        features = self.forward('fpn', inputs)
        export(self, ('fpn', inputs),
               export_name,
               input_names=['images'],
               output_names=list(features.keys()),
               opset_version=11)
        return features

    def export_fcoshead(self, features, export_name, input_names=None, output_names=None):
        self.precal_locations = self.proposal_generator.compute_locations(features)
        proposals = self.forward('fcoshead', features)
        export(self, ('fcoshead', features),
               export_name,
               input_names=input_names,
               output_names=output_names,
               opset_version=11)
        return proposals

    def export_roihead(self, input_shape, export_name):
        inputs = torch.randn(*input_shape)
        mask = self.forward('centerhead', inputs)
        export(self, ('centerhead', inputs), export_name,
               input_names=['roipool'],
               output_names=['mask'],
               opset_version=11)
        return mask

    def export(self, cfg, export_dir='.'):
        feat = self.export_fpn(
            (1, 3, 128, 128),
            f'{export_dir}/centermask_fpn.onnx')
        feat = [feat[k] for k in cfg.MODEL.FCOS.IN_FEATURES]
        scores, detections, centerness = self.export_fcoshead(
            feat,
            f'{export_dir}/centermask_fcos.onnx',
            cfg.MODEL.FCOS.IN_FEATURES,
            ['logit', 'box', 'centerness'])
        mask_inc = cfg.MODEL.ROI_MASK_HEAD.CONV_DIM
        mask_size = cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION
        self.export_roihead(
            (1, mask_inc, mask_size, mask_size),
            f'{export_dir}/centermask_roi.onnx')


def main(args):
    cfg = setup(args)
    model = RCNNExporter(cfg)
    DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
        cfg.MODEL.WEIGHTS, resume=args.resume
    )
    model = model.eval().to('cpu')
    model.export(cfg, cfg.OUTPUT_DIR)


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    main(args)

import cv2 as cv
import numpy as np
import torch
from detectron2.modeling import GeneralizedRCNN
from detectron2.layers import batched_nms
from torch.onnx import export

from train_net import DetectionCheckpointer, default_argument_parser, setup


class RCNNExporter(GeneralizedRCNN):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def preprocess_image(self, x):
        """
        Normalize, pad and batch the input images.
        """
        x -= self.pixel_mean.view(-1, 3, 1, 1)
        x /= self.pixel_std.view(-1, 3, 1, 1)
        return x

    def forward(self, inference_type, *x):
        if inference_type == 'fpn':
            inputs = self.preprocess_image(*x)
            features = self.backbone(inputs)
            return features
        elif inference_type == 'fcoshead':
            logits_pred, reg_pred, ctrness_pred, _ = self.proposal_generator.fcos_head(*x)
            strides = self.proposal_generator.strides
            bundle = (self.precal_locations, logits_pred, reg_pred, ctrness_pred, strides)
            scores = []
            centerness = []
            detections = []
            for l, o, r, c, s in zip(*bundle):
                r *= s
                N, C, H, W = o.shape
                box_cls = o.permute(0, 2, 3, 1).reshape(N, -1, C).sigmoid()
                box_regression = r.view(N, 4, H, W).permute(0, 2, 3, 1).reshape(N, -1, 4)
                ctrness = c.view(N, 1, H, W).permute(0, 2, 3, 1).reshape(N, -1).sigmoid()
                detection = torch.stack([
                    l[:, 0][None, ...] - box_regression[..., 0],
                    l[:, 1][None, ...] - box_regression[..., 1],
                    l[:, 0][None, ...] + box_regression[..., 2],
                    l[:, 1][None, ...] + box_regression[..., 3],
                ], dim=-1)
                scores.append(box_cls)
                centerness.append(ctrness)
                detections.append(detection)
            scores = torch.cat(scores, 1)
            centerness = torch.cat(centerness, 1)
            detections = torch.cat(detections, 1)
            return scores, detections / 128, centerness
        elif inference_type == 'centerhead':
            mask = self.roi_heads.mask_head(*x).sigmoid()
            return mask

    def export_fpn(self, input_shape, export_name):
        inputs = torch.randn(*input_shape)
        features = self.forward('fpn', inputs)
        export(self, ('fpn', inputs),
               export_name,
               input_names=['images'],
               output_names=list(features.keys()),
               opset_version=11)
        return features

    def export_fcoshead(self, features, export_name, input_names=None, output_names=None):
        self.precal_locations = self.proposal_generator.compute_locations(features)
        proposals = self.forward('fcoshead', features)
        export(self, ('fcoshead', features),
               export_name,
               input_names=input_names,
               output_names=output_names,
               opset_version=11)
        return proposals

    def export_roihead(self, input_shape, export_name):
        inputs = torch.randn(*input_shape)
        mask = self.forward('centerhead', inputs)
        export(self, ('centerhead', inputs), export_name,
               input_names=['roipool'],
               output_names=['mask'],
               opset_version=11)
        return mask

    def export(self, cfg, export_dir='.'):
        feat = self.export_fpn(
            (1, 3, 128, 128),
            f'{export_dir}/centermask_fpn.onnx')
        feat = [feat[k] for k in cfg.MODEL.FCOS.IN_FEATURES]
        scores, detections, centerness = self.export_fcoshead(
            feat,
            f'{export_dir}/centermask_fcos.onnx',
            cfg.MODEL.FCOS.IN_FEATURES,
            ['logit', 'box', 'centerness'])
        mask_inc = cfg.MODEL.ROI_MASK_HEAD.CONV_DIM
        mask_size = cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION
        self.export_roihead(
            (1, mask_inc, mask_size, mask_size),
            f'{export_dir}/centermask_roi.onnx')


def main(args):
    cfg = setup(args)
    model = RCNNExporter(cfg)
    DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
        cfg.MODEL.WEIGHTS, resume=args.resume
    )
    model = model.eval().to('cpu')
    model.export(cfg, cfg.OUTPUT_DIR)


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    main(args)
Normalize, pad and batch the input images.