Spaces:
Runtime error
Runtime error
fifa-tryon-demo
/
Self-Correction-Human-Parsing-for-ACGPN
/mhp_extension
/detectron2
/tests
/test_export_caffe2.py
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| # -*- coding: utf-8 -*- | |
| import copy | |
| import numpy as np | |
| import os | |
| import tempfile | |
| import unittest | |
| import cv2 | |
| import torch | |
| from fvcore.common.file_io import PathManager | |
| from detectron2 import model_zoo | |
| from detectron2.checkpoint import DetectionCheckpointer | |
| from detectron2.config import get_cfg | |
| from detectron2.data import DatasetCatalog | |
| from detectron2.modeling import build_model | |
| from detectron2.utils.logger import setup_logger | |
| class TestCaffe2Export(unittest.TestCase): | |
| def setUp(self): | |
| setup_logger() | |
| def _test_model(self, config_path, device="cpu"): | |
| # requires extra dependencies | |
| from detectron2.export import Caffe2Model, add_export_config, export_caffe2_model | |
| cfg = get_cfg() | |
| cfg.merge_from_file(model_zoo.get_config_file(config_path)) | |
| cfg = add_export_config(cfg) | |
| cfg.MODEL.DEVICE = device | |
| model = build_model(cfg) | |
| DetectionCheckpointer(model).load(model_zoo.get_checkpoint_url(config_path)) | |
| inputs = [{"image": self._get_test_image()}] | |
| c2_model = export_caffe2_model(cfg, model, copy.deepcopy(inputs)) | |
| with tempfile.TemporaryDirectory(prefix="detectron2_unittest") as d: | |
| c2_model.save_protobuf(d) | |
| c2_model.save_graph(os.path.join(d, "test.svg"), inputs=copy.deepcopy(inputs)) | |
| c2_model = Caffe2Model.load_protobuf(d) | |
| c2_model(inputs)[0]["instances"] | |
| def _get_test_image(self): | |
| try: | |
| file_name = DatasetCatalog.get("coco_2017_train")[0]["file_name"] | |
| assert PathManager.exists(file_name) | |
| except Exception: | |
| self.skipTest("COCO dataset not available.") | |
| with PathManager.open(file_name, "rb") as f: | |
| buf = f.read() | |
| img = cv2.imdecode(np.frombuffer(buf, dtype=np.uint8), cv2.IMREAD_COLOR) | |
| assert img is not None, file_name | |
| return torch.from_numpy(img.transpose(2, 0, 1)) | |
| def testMaskRCNN(self): | |
| self._test_model("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") | |
| def testMaskRCNNGPU(self): | |
| self._test_model("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml", device="cuda") | |
| def testRetinaNet(self): | |
| self._test_model("COCO-Detection/retinanet_R_50_FPN_3x.yaml") | |
| def testPanopticFPN(self): | |
| self._test_model("COCO-PanopticSegmentation/panoptic_fpn_R_50_3x.yaml") | |