@@ -0,0 +1,2634 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved%0A# -*- coding: utf-8 -*-%0A%0Aimport copy%0Aimport numpy as np%0Aimport os%0Aimport tempfile%0Aimport unittest%0Aimport cv2%0Aimport torch%0Afrom fvcore.common.file_io import PathManager%0A%0Afrom detectron2 import model_zoo%0Afrom detectron2.checkpoint import DetectionCheckpointer%0Afrom detectron2.config import get_cfg%0Afrom detectron2.data import DatasetCatalog%0Afrom detectron2.modeling import build_model%0Afrom detectron2.utils.logger import setup_logger%0A%0A%0Aclass TestCaffe2Export(unittest.TestCase):%0A    def setUp(self):%0A        setup_logger()%0A        if os.environ.get(%22CIRCLECI%22):%0A            unittest.skipTest(%22Require COCO data and model zoo.%22)%0A%0A    def _test_model(self, config_path, device=%22cpu%22):%0A        # requires extra dependencies%0A        from detectron2.export import Caffe2Model, add_export_config, export_caffe2_model%0A%0A        cfg = get_cfg()%0A        cfg.merge_from_file(model_zoo.get_config_file(config_path))%0A        cfg = add_export_config(cfg)%0A        cfg.MODEL.DEVICE = device%0A%0A        model = build_model(cfg)%0A        DetectionCheckpointer(model).load(model_zoo.get_checkpoint_url(config_path))%0A%0A        inputs = %5B%7B%22image%22: self._get_test_image()%7D%5D%0A        c2_model = export_caffe2_model(cfg, model, copy.deepcopy(inputs))%0A%0A        with tempfile.TemporaryDirectory(prefix=%22detectron2_unittest%22) as d:%0A            c2_model.save_protobuf(d)%0A            c2_model.save_graph(os.path.join(d, %22test.svg%22), inputs=copy.deepcopy(inputs))%0A            c2_model = Caffe2Model.load_protobuf(d)%0A        c2_model(inputs)%5B0%5D%5B%22instances%22%5D%0A%0A    def _get_test_image(self):%0A        try:%0A            file_name = DatasetCatalog.get(%22coco_2017_train%22)%5B0%5D%5B%22file_name%22%5D%0A            assert PathManager.exists(file_name)%0A        except Exception:%0A            unittest.SkipTest(%22COCO dataset not available.%22)%0A%0A        with PathManager.open(file_name, %22rb%22) as f:%0A            buf = f.read()%0A        img = cv2.imdecode(np.frombuffer(buf, dtype=np.uint8), cv2.IMREAD_COLOR)%0A        assert img is not None, file_name%0A        return torch.from_numpy(img.transpose(2, 0, 1))%0A%0A    def testMaskRCNN(self):%0A        self._test_model(%22COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml%22)%0A%0A    @unittest.skipIf(not torch.cuda.is_available(), %22CUDA not available%22)%0A    def testMaskRCNNGPU(self):%0A        self._test_model(%22COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml%22, device=%22cuda%22)%0A%0A    def testRetinaNet(self):%0A        self._test_model(%22COCO-Detection/retinanet_R_50_FPN_3x.yaml%22)%0A%0A    def testPanopticFPN(self):%0A        self._test_model(%22COCO-PanopticSegmentation/panoptic_fpn_R_50_3x.yaml%22)%0A

make unittest "test_export" public