Initial commit

fbshipit-source-id: f4a8ba78691d8cf46e003ef0bd2e95f170932778

tests/resources/arch_def_merging.yaml

+MODEL:
+  FBNET_V2:
+    ARCH: ""
+    ARCH_DEF: [{
+      "trunk": [
+          [["conv_k3", 8, 2, 1], ["conv_k3", 8, 1, 1, {"expansion": 1}]],
+          [["conv_k3", 10, 2, 1, {"expansion": 4}]],
+          [["conv_k3", 24, 2, 2, {"expansion": 6}]],
+          [["conv_k3", 40, 2, 2, {"expansion": 6}], ["conv_k3", 64, 1, 2, {"expansion": 6}]]
+      ],
+      "kpts": [
+          [["conv_k3", 112, 2, 3, {"expansion": 6}], ["conv_k3", 160, 1, 1, {"expansion": 4}]]
+      ],
+      "bbox": [
+          [["conv_k3", 112, 2, 3, {"expansion": 6}], ["conv_k3", 160, 1, 1, {"expansion": 4}]]
+      ],
+      "rpn": [
+          [["conv_k3", 64, 1, 1, {"expansion": 6}]]
+      ]
+    }]

tests/test_box2box_transform.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import numpy as np
+import torch
+import unittest
+
+from detectron2.modeling.box_regression import Box2BoxTransform
+
+
+class TestBox2BoxTransform(unittest.TestCase):
+    def test_box2box_transform(self):
+        """ Match unit test UtilsBoxesTest.TestBboxTransformRandom in
+            caffe2/operators/generate_proposals_op_util_boxes_test.cc
+        """
+        box2box_transform = Box2BoxTransform(weights=(1.0, 1.0, 1.0, 1.0))
+        bbox = torch.from_numpy(
+            np.array(
+                [
+                    175.62031555,
+                    20.91103172,
+                    253.352005,
+                    155.0145874,
+                    169.24636841,
+                    4.85241556,
+                    228.8605957,
+                    105.02092743,
+                    181.77426147,
+                    199.82876587,
+                    192.88427734,
+                    214.0255127,
+                    174.36262512,
+                    186.75761414,
+                    296.19091797,
+                    231.27906799,
+                    22.73153877,
+                    92.02596283,
+                    135.5695343,
+                    208.80291748,
+                ]
+            )
+            .astype(np.float32)
+            .reshape(-1, 4)
+        )
+
+        deltas = torch.from_numpy(
+            np.array(
+                [
+                    0.47861834,
+                    0.13992102,
+                    0.14961673,
+                    0.71495209,
+                    0.29915856,
+                    -0.35664671,
+                    0.89018666,
+                    0.70815367,
+                    -0.03852064,
+                    0.44466892,
+                    0.49492538,
+                    0.71409376,
+                    0.28052918,
+                    0.02184832,
+                    0.65289006,
+                    1.05060139,
+                    -0.38172557,
+                    -0.08533806,
+                    -0.60335309,
+                    0.79052375,
+                ]
+            )
+            .astype(np.float32)
+            .reshape(-1, 4)
+        )
+
+        gt_bbox = (
+            np.array(
+                [
+                    206.949539,
+                    -30.715202,
+                    297.387665,
+                    244.448486,
+                    143.871216,
+                    -83.342888,
+                    290.502289,
+                    121.053398,
+                    177.430283,
+                    198.666245,
+                    196.295273,
+                    228.703079,
+                    152.251892,
+                    145.431564,
+                    387.215454,
+                    274.594238,
+                    5.062420,
+                    11.040955,
+                    66.328903,
+                    269.686218,
+                ]
+            )
+            .astype(np.float32)
+            .reshape(-1, 4)
+        )
+
+        # Detectron2 removed box plus one
+        bbox[:, 2] += 1
+        bbox[:, 3] += 1
+        gt_bbox[:, 2] += 1
+        gt_bbox[:, 3] += 1
+
+        results = box2box_transform.apply_deltas(deltas, bbox)
+
+        np.testing.assert_allclose(results.detach().numpy(), gt_bbox, atol=1e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_box_with_nms_limit.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import torch
+import unittest
+
+from detectron2.modeling.roi_heads.fast_rcnn import fast_rcnn_inference
+from detectron2.layers import cat
+from detectron2.structures import Boxes
+
+
+class TestBoxWithNMSLimit(unittest.TestCase):
+    def test_caffe2_pytorch_eq(self):
+        ims_per_batch = 8
+        post_nms_topk = 100
+        detections_per_im = 10
+        num_class = 80
+        score_thresh = 0.05
+        nms_thresh = 0.5
+
+        image_shapes = [torch.Size([800, 800])] * ims_per_batch
+        batch_splits = [post_nms_topk] * ims_per_batch
+
+        # NOTE: There're still some unsure minor implementation differences
+        # (eg. ordering when equal score across classes) causing some seeds
+        # don't pass the test.
+        # Thus set a fixed seed to make sure this test passes consistantly.
+        rng = torch.Generator()
+        rng.manual_seed(42)
+        boxes = []
+        for n in batch_splits:
+            box = 1000.0 * 0.5 * torch.rand(n, num_class, 4, generator=rng) + 0.001
+            box[:, :, -2:] += box[:, :, :2]
+            box = box.view(n, num_class * 4)
+            boxes.append(box)
+        scores = [torch.rand(n, num_class + 1, generator=rng) for n in batch_splits]
+
+        ref_results, ref_kept_indices = fast_rcnn_inference(
+            boxes, scores, image_shapes,
+            score_thresh=score_thresh,
+            nms_thresh=nms_thresh,
+            topk_per_image=detections_per_im
+        )
+        for result, kept_index, score in zip(ref_results, ref_kept_indices, scores):
+            torch.testing.assert_allclose(
+                score[kept_index, result.pred_classes],
+                result.scores,
+            )
+
+        # clip is done in BBoxTransformOp
+        c2_boxes = []
+        for box, image_shape in zip(boxes, image_shapes):
+            num_bbox_reg_classes = box.shape[1] // 4
+            clipped_box = Boxes(box.reshape(-1, 4))
+            clipped_box.clip(image_shape)
+            clipped_box = clipped_box.tensor.view(-1, num_bbox_reg_classes * 4)
+            c2_boxes.append(clipped_box)
+
+        c2_boxes = cat(c2_boxes)
+        c2_scores = cat(scores)
+        c2_batch_splits = torch.Tensor(batch_splits)
+
+        nms_outputs = torch.ops._caffe2.BoxWithNMSLimit(
+            c2_scores,
+            c2_boxes,
+            c2_batch_splits,
+            score_thresh=float(score_thresh),
+            nms=float(nms_thresh),
+            detections_per_im=int(detections_per_im),
+            soft_nms_enabled=False,
+            soft_nms_method="linear",
+            soft_nms_sigma=0.5,
+            soft_nms_min_score_thres=0.001,
+            rotated=False,
+            cls_agnostic_bbox_reg=False,
+            input_boxes_include_bg_cls=False,
+            output_classes_include_bg_cls=False,
+            legacy_plus_one=False,
+        )
+        roi_score_nms, roi_bbox_nms, roi_class_nms, roi_batch_splits_nms, roi_keeps_nms, roi_keeps_size_nms = nms_outputs  # noqa
+
+        roi_score_nms = roi_score_nms.split(roi_batch_splits_nms.int().tolist())
+        roi_bbox_nms = roi_bbox_nms.split(roi_batch_splits_nms.int().tolist())
+        roi_class_nms = roi_class_nms.split(roi_batch_splits_nms.int().tolist())
+        roi_keeps_nms = roi_keeps_nms.split(roi_batch_splits_nms.int().tolist())
+
+        for _score_nms, _class_nms, _keeps_nms, _score in zip(
+            roi_score_nms, roi_class_nms, roi_keeps_nms, scores
+        ):
+            torch.testing.assert_allclose(
+                _score[_keeps_nms.to(torch.int64), _class_nms.to(torch.int64)],
+                _score_nms,
+            )
+
+        for ref, s, b, c in zip(
+            ref_results, roi_score_nms, roi_bbox_nms, roi_class_nms
+        ):
+            s1, i1 = s.sort()
+            s2, i2 = ref.scores.sort()
+            torch.testing.assert_allclose(s1, s2)
+            torch.testing.assert_allclose(b[i1], ref.pred_boxes.tensor[i2])
+            torch.testing.assert_allclose(c.to(torch.int64)[i1], ref.pred_classes[i2])
+
+        for ref, k in zip(ref_kept_indices, roi_keeps_nms):
+            # NOTE: order might be different due to implementation
+            ref_set = set(ref.tolist())
+            k_set = set(k.tolist())
+            self.assertEqual(ref_set, k_set)

tests/test_config_utils.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import copy
+import unittest
+
+from d2go.config.utils import (
+    config_dict_to_list_str,
+    flatten_config_dict,
+    str_wrap_fbnet_arch_def,
+)
+
+
+class TestConfigUtils(unittest.TestCase):
+    def test_str_wrap_fbnet_arch_def(self):
+        """Check that fbnet modeldef converted to str"""
+        d = {"MODEL": {"FBNET_V2": {"ARCH_DEF": {"key0": "val0"}}}}
+        new_dict = str_wrap_fbnet_arch_def(d)
+        gt = {"MODEL": {"FBNET_V2": {"ARCH_DEF": """'{"key0": "val0"}'"""}}}
+        self.assertEqual(new_dict, gt)
+        self.assertNotEqual(d, new_dict)
+
+        # check only fbnet arch is changed
+        d = {"a0": "a1", "b0": {"b1": "b2"}}
+        gt = copy.deepcopy(d)
+        new_dict = str_wrap_fbnet_arch_def(d)
+        self.assertEqual(new_dict, gt)
+
+    def test_flatten_config_dict(self):
+        """Check flatten config dict to single layer dict"""
+        d = {"a0": "a1", "b0": {"b1": "b2"}, "c0": {"c1": {"c2": 3}}}
+        fdict = flatten_config_dict(d)
+        gt = {"a0": "a1", "b0.b1": "b2", "c0.c1.c2": 3}
+        self.assertEqual(fdict, gt)
+
+    def test_config_dict_to_list_str(self):
+        """Check convert config dict to str list"""
+        d = {"a0": "a1", "b0": {"b1": "b2"}, "c0": {"c1": {"c2": 3}}}
+        str_list = config_dict_to_list_str(d)
+        gt = ["a0", "a1", "b0.b1", "b2", "c0.c1.c2", "3"]
+        self.assertEqual(str_list, gt)

tests/test_configs.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import glob
+import logging
+import os
+import unittest
+
+from d2go.config import auto_scale_world_size, reroute_config_path
+from d2go.runner import GeneralizedRCNNRunner
+from mobile_cv.common.misc.file_utils import make_temp_directory
+
+
+logger = logging.getLogger(__name__)
+
+
+class TestConfigs(unittest.TestCase):
+    def test_configs_load(self):
+        """ Make sure configs are loadable """
+
+        for location in ["detectron2", "detectron2go"]:
+            root_dir = os.path.abspath(reroute_config_path(f"{location}://."))
+            files = glob.glob(os.path.join(root_dir, "**/*.yaml"), recursive=True)
+            self.assertGreater(len(files), 0)
+            for fn in sorted(files):
+                logger.info("Loading {}...".format(fn))
+                GeneralizedRCNNRunner().get_default_cfg().merge_from_file(fn)
+
+    def test_arch_def_loads(self):
+        """ Test arch def str-to-dict conversion compatible with merging """
+        default_cfg = GeneralizedRCNNRunner().get_default_cfg()
+        cfg = default_cfg.clone()
+        cfg.merge_from_file(os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                         "resources/arch_def_merging.yaml"))
+
+        with make_temp_directory("detectron2go_tmp") as tmp_dir:
+            # Dump out config with arch def
+            file_name = os.path.join(tmp_dir, "test_archdef_config.yaml")
+            with open(file_name, "w") as f:
+                f.write(cfg.dump())
+
+            # Attempt to reload the config
+            another_cfg = default_cfg.clone()
+            another_cfg.merge_from_file(file_name)
+
+    def test_default_cfg_dump_and_load(self):
+        default_cfg = GeneralizedRCNNRunner().get_default_cfg()
+
+        cfg = default_cfg.clone()
+        with make_temp_directory("detectron2go_tmp") as tmp_dir:
+            file_name = os.path.join(tmp_dir, "config.yaml")
+            # this is same as the one in fblearner_launch_utils_detectron2go.py
+            with open(file_name, "w") as f:
+                f.write(cfg.dump(default_flow_style=False))
+
+            # check if the dumped config file can be merged
+            cfg.merge_from_file(file_name)
+
+    def test_default_cfg_deprecated_keys(self):
+        default_cfg = GeneralizedRCNNRunner().get_default_cfg()
+
+        # a warning will be printed for deprecated keys
+        default_cfg.merge_from_list(["QUANTIZATION.QAT.LOAD_PRETRAINED", True])
+        # exception will raise for renamed keys
+        self.assertRaises(
+            KeyError,
+            default_cfg.merge_from_list,
+            ["QUANTIZATION.QAT.BACKEND", "fbgemm"],
+        )
+
+
+class TestAutoScaleWorldSize(unittest.TestCase):
+    def test_8gpu_to_1gpu(self):
+        """
+        when scaling a 8-gpu config to 1-gpu one, the batch size will be reduced by 8x
+        """
+        cfg = GeneralizedRCNNRunner().get_default_cfg()
+        self.assertEqual(cfg.SOLVER.REFERENCE_WORLD_SIZE, 8)
+        batch_size_x8 = cfg.SOLVER.IMS_PER_BATCH
+        assert batch_size_x8 % 8 == 0, "default batch size is not multiple of 8"
+        auto_scale_world_size(cfg, new_world_size=1)
+        self.assertEqual(cfg.SOLVER.REFERENCE_WORLD_SIZE, 1)
+        self.assertEqual(cfg.SOLVER.IMS_PER_BATCH * 8, batch_size_x8)
+
+    def test_not_scale_for_zero_world_size(self):
+        """
+        when reference world size is 0, no scaling should happen
+        """
+        cfg = GeneralizedRCNNRunner().get_default_cfg()
+        self.assertEqual(cfg.SOLVER.REFERENCE_WORLD_SIZE, 8)
+        cfg.SOLVER.REFERENCE_WORLD_SIZE = 0
+        batch_size_x8 = cfg.SOLVER.IMS_PER_BATCH
+        auto_scale_world_size(cfg, new_world_size=1)
+        self.assertEqual(cfg.SOLVER.REFERENCE_WORLD_SIZE, 0)
+        self.assertEqual(cfg.SOLVER.IMS_PER_BATCH, batch_size_x8)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_d2go_datasets.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import json
+import os
+import unittest
+
+from detectron2.data import DatasetCatalog
+from d2go.data.utils import maybe_subsample_n_images
+from d2go.runner import Detectron2GoRunner
+from mobile_cv.common.misc.file_utils import make_temp_directory
+
+from d2go.tests.data_loader_helper import LocalImageGenerator, create_toy_dataset
+
+
+def create_test_images_and_dataset_json(data_dir):
+    # create image and json
+    image_dir = os.path.join(data_dir, "images")
+    os.makedirs(image_dir)
+    json_dataset, meta_data = create_toy_dataset(
+        LocalImageGenerator(image_dir, width=80, height=60), num_images=10
+    )
+    json_file = os.path.join(data_dir, "{}.json".format("inj_ds1"))
+    with open(json_file, "w") as f:
+        json.dump(json_dataset, f)
+
+    return image_dir, json_file
+
+
+class TestD2GoDatasets(unittest.TestCase):
+    def test_coco_injection(self):
+
+        with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir:
+            image_dir, json_file = create_test_images_and_dataset_json(tmp_dir)
+
+            runner = Detectron2GoRunner()
+            cfg = runner.get_default_cfg()
+            cfg.merge_from_list(
+                [
+                    str(x)
+                    for x in [
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.NAMES",
+                        ["inj_ds1", "inj_ds2"],
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.IM_DIRS",
+                        [image_dir, "/mnt/fair"],
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.JSON_FILES",
+                        [json_file, "inj_ds2"],
+                    ]
+                ]
+            )
+
+            runner.register(cfg)
+            inj_ds1 = DatasetCatalog.get("inj_ds1")
+            self.assertEqual(len(inj_ds1), 10)
+            for dic in inj_ds1:
+                self.assertEqual(dic["width"], 80)
+                self.assertEqual(dic["height"], 60)
+
+    def test_sub_dataset(self):
+        with make_temp_directory("detectron2go_tmp_dataset") as tmp_dir:
+            image_dir, json_file = create_test_images_and_dataset_json(tmp_dir)
+
+            runner = Detectron2GoRunner()
+            cfg = runner.get_default_cfg()
+            cfg.merge_from_list(
+                [
+                    str(x)
+                    for x in [
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.NAMES",
+                        ["inj_ds"],
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.IM_DIRS",
+                        [image_dir],
+                        "D2GO_DATA.DATASETS.COCO_INJECTION.JSON_FILES",
+                        [json_file],
+                        "DATASETS.TEST",
+                        ("inj_ds",),
+                        "D2GO_DATA.TEST.MAX_IMAGES",
+                        1,
+                    ]
+                ]
+            )
+
+            runner.register(cfg)
+            with maybe_subsample_n_images(cfg) as new_cfg:
+                test_loader = runner.build_detection_test_loader(
+                    new_cfg, new_cfg.DATASETS.TEST[0]
+                )
+                self.assertEqual(len(test_loader), 1)

tests/test_data_loader.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import os
+import unittest
+
+from d2go.runner import GeneralizedRCNNRunner, create_runner
+from mobile_cv.common.misc.file_utils import make_temp_directory
+from PIL import Image
+
+from d2go.tests.data_loader_helper import LocalImageGenerator, register_toy_dataset
+
+
+class TestD2GoDatasetMapper(unittest.TestCase):
+    """
+    This class test D2GoDatasetMapper which is used to build
+    data loader in GeneralizedRCNNRunner (the default runner) in Detectron2Go.
+    """
+
+    def test_default_dataset(self):
+        runner = create_runner("d2go.runner.GeneralizedRCNNRunner")
+        cfg = runner.get_default_cfg()
+        cfg.DATASETS.TRAIN = ["default_dataset_train"]
+        cfg.DATASETS.TEST = ["default_dataset_test"]
+
+        with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir:
+            image_dir = os.path.join(dataset_dir, "images")
+            os.makedirs(image_dir)
+            image_generator = LocalImageGenerator(image_dir, width=80, height=60)
+
+            with register_toy_dataset(
+                "default_dataset_train", image_generator, num_images=3
+            ):
+                train_loader = runner.build_detection_train_loader(cfg)
+                for i, data in enumerate(train_loader):
+                    self.assertIsNotNone(data)
+                    # for training loader, it has infinite length
+                    if i == 6:
+                        break
+
+            with register_toy_dataset(
+                "default_dataset_test", image_generator, num_images=3
+            ):
+                test_loader = runner.build_detection_test_loader(
+                    cfg, dataset_name="default_dataset_test"
+                )
+                all_data = []
+                for data in test_loader:
+                    all_data.append(data)
+                self.assertEqual(len(all_data), 3)

tests/test_data_transforms.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import numpy as np
+from detectron2.data.transforms import apply_transform_gens
+from d2go.data.transforms.build import build_transform_gen
+from d2go.runner import Detectron2GoRunner
+
+
+class TestDataTransforms(unittest.TestCase):
+    def test_build_transform_gen(self):
+        default_cfg = Detectron2GoRunner().get_default_cfg()
+        default_cfg.INPUT.MIN_SIZE_TRAIN = (30,)
+        default_cfg.INPUT.MIN_SIZE_TEST = 30
+
+        trans_train = build_transform_gen(default_cfg, is_train=True)
+        trans_test = build_transform_gen(default_cfg, is_train=False)
+
+        img = np.zeros((80, 60, 3))
+        trans_img_train, tl_train = apply_transform_gens(trans_train, img)
+        trans_img_test, tl_test = apply_transform_gens(trans_test, img)
+
+        self.assertEqual(trans_img_train.shape, (40, 30, 3))
+        self.assertEqual(trans_img_test.shape, (40, 30, 3))
+
+    def test_build_transform_gen_resize_square(self):
+        default_cfg = Detectron2GoRunner().get_default_cfg()
+        default_cfg.INPUT.MIN_SIZE_TRAIN = (30,)
+        default_cfg.INPUT.MIN_SIZE_TEST = 40
+        default_cfg.D2GO_DATA.AUG_OPS.TRAIN = ["ResizeShortestEdgeSquareOp"]
+        default_cfg.D2GO_DATA.AUG_OPS.TEST = ["ResizeShortestEdgeSquareOp"]
+
+        trans_train = build_transform_gen(default_cfg, is_train=True)
+        trans_test = build_transform_gen(default_cfg, is_train=False)
+
+        img = np.zeros((80, 60, 3))
+        trans_img_train, tl_train = apply_transform_gens(trans_train, img)
+        trans_img_test, tl_test = apply_transform_gens(trans_test, img)
+
+        self.assertEqual(trans_img_train.shape, (30, 30, 3))
+        self.assertEqual(trans_img_test.shape, (40, 40, 3))

tests/test_data_transforms_box_utils.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import d2go.data.transforms.box_utils as bu
+import numpy as np
+import torch
+
+
+class TestDataTransformsBoxUtils(unittest.TestCase):
+    def test_min_box_ar(self):
+        box_xywh = [4, 5, 10, 6]
+        target_aspect_ratio = 1.0 / 2
+        new_box = bu.get_min_box_aspect_ratio(box_xywh, target_aspect_ratio)
+        self.assertArrayEqual(torch.Tensor([4, -2, 10, 20]), new_box)
+
+    def test_get_box_from_mask(self):
+        img_w, img_h = 8, 6
+        mask = np.zeros([img_h, img_w])
+        self.assertEqual(mask.shape, (img_h, img_w))
+        mask[2:4, 3:6] = 1
+        box = bu.get_box_from_mask(mask)
+        self.assertEqual(box, (3, 2, 3, 2))
+
+    def test_get_box_from_mask_union(self):
+        img_w, img_h = 8, 6
+        mask = np.zeros([img_h, img_w])
+        self.assertEqual(mask.shape, (img_h, img_w))
+        mask[2:4, 1:4] = 1
+        mask[5:6, 4:8] = 1
+        box = bu.get_box_from_mask(mask)
+        self.assertEqual(box, (1, 2, 7, 4))
+
+    def test_get_box_from_mask_empty(self):
+        img_w, img_h = 8, 6
+        mask = np.zeros([img_h, img_w])
+        box = bu.get_box_from_mask(mask)
+        self.assertIsNone(box)
+
+    def test_scale_bbox_center(self):
+        bbox = torch.Tensor([1, 2, 4, 5])
+        out_bbox = bu.scale_bbox_center(bu.scale_bbox_center(bbox, 2.0), 0.5)
+        self.assertArrayEqual(bbox, out_bbox)
+
+    def assertArrayEqual(self, a1, a2):
+        self.assertTrue(np.array_equal(a1, a2))

tests/test_data_transforms_color_yuv.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import numpy as np
+from d2go.data.transforms import color_yuv as cy
+from d2go.runner import Detectron2GoRunner
+from detectron2.data.transforms import apply_augmentations
+from d2go.data.transforms.build import build_transform_gen
+
+
+class TestDataTransformsColorYUV(unittest.TestCase):
+
+    def test_yuv_color_transforms(self):
+        default_cfg = Detectron2GoRunner().get_default_cfg()
+        img = np.concatenate([
+            np.random.uniform(0, 1, size=(80, 60, 1)),
+            np.random.uniform(-0.5, 0.5, size=(80, 60, 1)),
+            np.random.uniform(-0.5, 0.5, size=(80, 60, 1)),
+        ], axis=2)
+
+        default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [
+            'RandomContrastYUVOp::{"intensity_min": 0.3, "intensity_max": 0.5}',
+        ]
+        low_contrast_tfm = build_transform_gen(default_cfg, is_train=True)
+        low_contrast, _ = apply_augmentations(low_contrast_tfm, img)
+
+        default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [
+            'RandomSaturationYUVOp::{"intensity_min": 1.5, "intensity_max": 1.7}',
+        ]
+        high_saturation_tfm = build_transform_gen(default_cfg, is_train=True)
+        high_saturation, _ = apply_augmentations(high_saturation_tfm, img)
+
+        # Use pixel statistics to roughly check transformed images as expected
+        # All channels have less variance
+        self.assertLess(np.var(low_contrast[:, :, 0]), np.var(img[:, :, 0]))
+        self.assertLess(np.var(low_contrast[:, :, 1]), np.var(img[:, :, 1]))
+        self.assertLess(np.var(low_contrast[:, :, 2]), np.var(img[:, :, 2]))
+
+        # 1st channel is unchanged (test w/ mean, var), 2nd + 3rd channels more variance
+        self.assertAlmostEqual(np.mean(high_saturation[:, :, 0]), np.mean(img[:, :, 0]))
+        self.assertAlmostEqual(np.var(high_saturation[:, :, 0]), np.var(img[:, :, 0]))
+        self.assertGreater(np.var(high_saturation[:, :, 1]), np.var(img[:, :, 1]))
+        self.assertGreater(np.var(high_saturation[:, :, 2]), np.var(img[:, :, 2]))
+
+
+    def test_transform_color_yuv_rgbyuv_convert(self):
+        image = np.arange(256).reshape(16, 16, 1).repeat(3, axis=2).astype(np.uint8)
+        tf1 = cy.RGB2YUVBT601().get_transform(image)
+        tf2 = cy.YUVBT6012RGB().get_transform(image)
+
+        image_yuv = tf1.apply_image(image)
+        image_rgb = tf2.apply_image(image_yuv)
+
+        self.assertArrayEqual((image_rgb + 0.5).astype(np.uint8), image)
+
+    def test_transform_color_yuv_rgbyuv_convert_invese(self):
+        image = np.arange(256).reshape(16, 16, 1).repeat(3, axis=2).astype(np.uint8)
+
+        tf = cy.RGB2YUVBT601().get_transform(image)
+
+        image_yuv = tf.apply_image(image)
+        image_rgb = tf.inverse().apply_image(image_yuv)
+
+        self.assertArrayEqual((image_rgb + 0.5).astype(np.uint8), image)
+
+    def assertArrayEqual(self, a1, a2):
+        self.assertTrue(np.array_equal(a1, a2))

tests/test_data_transforms_crop.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import d2go.data.transforms.box_utils as bu
+import numpy as np
+import torch
+from d2go.data.transforms import crop as tf_crop
+
+
+class TestDataTransformsCrop(unittest.TestCase):
+    def test_transform_crop_extent_transform(self):
+        img_wh = (16, 11)
+
+        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
+        # h, w
+        sem_seg[5, 4] = 1
+        sem_seg[10, 13] = 1
+        sem_seg[5:11, 4:14] = 1
+
+        # src_rect: [x0, y0, x1, y1] in pixel coordinate, output_size: [h, w]
+        trans = tf_crop.ExtentTransform(src_rect=[4, 5, 14, 11], output_size=[6, 10])
+        out_mask = trans.apply_segmentation(sem_seg)
+        self.assertArrayEqual(out_mask.shape, torch.Tensor([6, 10]))
+        self.assertArrayEqual(np.unique(out_mask), torch.Tensor([1]))
+
+        trans = tf_crop.ExtentTransform(src_rect=[3, 4, 15, 11], output_size=[7, 12])
+        out_mask = trans.apply_segmentation(sem_seg)
+        self.assertArrayEqual(out_mask.shape, torch.Tensor([7, 12]))
+        self.assertArrayEqual(np.unique(out_mask), torch.Tensor([0, 1]))
+        self.assertArrayEqual(np.unique(out_mask[1:, 1:-1]), torch.Tensor([1]))
+        self.assertEqual(out_mask[:, 0].sum(), 0)
+        self.assertArrayEqual(out_mask[0, :].sum(), 0)
+        self.assertArrayEqual(out_mask[:, -1].sum(), 0)
+
+    def test_transform_crop_random_crop_fixed_aspect_ratio(self):
+        aug = tf_crop.RandomCropFixedAspectRatio([1.0 / 2])
+        img_wh = (16, 11)
+
+        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
+        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
+        sem_seg[5, 4] = 1
+        sem_seg[10, 13] = 1
+        mask_xywh = bu.get_box_from_mask(sem_seg)
+        self.assertArrayEqual(mask_xywh, torch.Tensor([4, 5, 10, 6]))
+
+        trans = aug.get_transform(img, sem_seg)
+        self.assertArrayEqual(trans.src_rect, torch.Tensor([4, -2, 14, 18]))
+        self.assertArrayEqual(trans.output_size, torch.Tensor([20, 10]))
+
+        out_img = trans.apply_image(img)
+        self.assertArrayEqual(out_img.shape, torch.Tensor([20, 10, 3]))
+
+        self.assertArrayEqual(np.unique(out_img[2:13, :, :]), torch.Tensor([1]))
+        self.assertArrayEqual(np.unique(out_img[0:2, :, :]), torch.Tensor([0]))
+        self.assertArrayEqual(np.unique(out_img[13:, :, :]), torch.Tensor([0]))
+
+        out_mask = trans.apply_segmentation(sem_seg)
+        self.assertArrayEqual(out_mask.shape, torch.Tensor([20, 10]))
+        self.assertEqual(out_mask[7, 0], 1)
+        self.assertEqual(out_mask[12, -1], 1)
+
+    def test_transform_crop_random_crop_fixed_aspect_ratio_scale_offset(self):
+        aug = tf_crop.RandomCropFixedAspectRatio(
+            [1.0 / 2], scale_range=[0.5, 0.5], offset_scale_range=[-0.5, -0.5]
+        )
+        img_wh = (16, 11)
+
+        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
+        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
+        sem_seg[5, 4] = 1
+        sem_seg[10, 13] = 1
+        sem_seg[5:11, 4:14] = 1
+        mask_xywh = bu.get_box_from_mask(sem_seg)
+        self.assertArrayEqual(mask_xywh, torch.Tensor([4, 5, 10, 6]))
+
+        trans = aug.get_transform(img, sem_seg)
+        self.assertArrayEqual(trans.src_rect, torch.Tensor([1.5, 0.0, 6.5, 10.0]))
+        self.assertArrayEqual(trans.output_size, torch.Tensor([10, 5]))
+
+        out_img = trans.apply_image(img)
+        self.assertArrayEqual(out_img.shape, torch.Tensor([10, 5, 3]))
+        self.assertEqual(np.unique(out_img), 1)
+
+        out_mask = trans.apply_segmentation(sem_seg)
+        self.assertArrayEqual(out_mask.shape, torch.Tensor([10, 5]))
+        self.assertEqual(np.unique(out_mask[6:, 3:]), 1)
+
+    def test_transform_crop_random_crop_fixed_aspect_ratio_empty_mask(self):
+        """The sem_mask is empty (the whole image is background)"""
+        aug = tf_crop.RandomCropFixedAspectRatio([1.0 / 2])
+        img_wh = (16, 11)
+
+        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
+        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
+
+        mask_xywh = bu.get_box_from_mask(sem_seg)
+        self.assertEqual(mask_xywh, None)
+
+        trans = aug.get_transform(img, sem_seg)
+        self.assertIsInstance(trans, tf_crop.NoOpTransform)
+
+        out_img = trans.apply_image(img)
+        self.assertArrayEqual(out_img.shape, img.shape)
+
+        out_mask = trans.apply_segmentation(sem_seg)
+        self.assertArrayEqual(out_mask.shape, sem_seg.shape)
+
+    def test_pad_transform(self):
+        crop_w, crop_h = 4, 3
+        full_w, full_h = 11, 9
+        crop_x, crop_y = 5, 6
+        trans = tf_crop.PadTransform(crop_x, crop_y, crop_w, crop_h, full_w, full_h)
+        img = np.ones([crop_h, crop_w])
+        trans_img = trans.apply_image(img)
+        self.assertArrayEqual(trans_img.shape, [full_h, full_w])
+        self.assertArrayEqual(np.unique(trans_img), [0, 1])
+        full_img_gt = np.zeros([full_h, full_w])
+        full_img_gt[crop_y : (crop_y + crop_h), crop_x : (crop_x + crop_w)] = 1
+        self.assertArrayEqual(full_img_gt, trans_img)
+
+    def test_crop_transform_inverse(self):
+        crop_w, crop_h = 4, 3
+        full_w, full_h = 11, 9
+        crop_x, crop_y = 5, 6
+        trans = tf_crop.InvertibleCropTransform(
+            crop_x, crop_y, crop_w, crop_h, full_w, full_h
+        )
+
+        full_img_gt = np.zeros([full_h, full_w])
+        full_img_gt[crop_y : (crop_y + crop_h), crop_x : (crop_x + crop_w)] = 1
+        crop_img_gt = np.ones([crop_h, crop_w])
+
+        self.assertArrayEqual(trans.apply_image(full_img_gt), crop_img_gt)
+        self.assertArrayEqual(trans.inverse().apply_image(crop_img_gt), full_img_gt)
+        self.assertArrayEqual(
+            trans.inverse().inverse().apply_image(full_img_gt), crop_img_gt
+        )
+
+    def test_pad_transform(self):
+        img_h, img_w = 10, 7
+        divisibility = 8
+
+        aug = tf_crop.PadBorderDivisible(divisibility)
+
+        img = np.ones([img_h, img_w, 3]) * 3
+        trans = aug.get_transform(img)
+        pad_img = trans.apply_image(img)
+        self.assertEqual(pad_img.shape, (16, 8, 3))
+        inverse_img = trans.inverse().apply_image(pad_img)
+        self.assertEqual(inverse_img.shape, (10, 7, 3))
+        self.assertArrayEqual(img, inverse_img)
+
+        mask = np.ones([img_h, img_w]) * 2
+        pad_mask = trans.apply_segmentation(mask)
+        self.assertEqual(pad_mask.shape, (16, 8))
+        inverse_mask = trans.inverse().apply_segmentation(pad_mask)
+        self.assertEqual(inverse_mask.shape, (10, 7))
+        self.assertArrayEqual(mask, inverse_mask)
+
+    def test_random_instance_crop(self):
+        from detectron2.data import detection_utils as du
+        from detectron2.data.transforms.augmentation import (
+            AugInput,
+            AugmentationList,
+        )
+        from detectron2.structures import BoxMode
+
+        aug = tf_crop.RandomInstanceCrop([1.0, 1.0])
+
+        img_w, img_h = 10, 7
+        annotations = [
+            {
+                "category_id": 0,
+                "bbox": [1, 1, 4, 3],
+                "bbox_mode": BoxMode.XYWH_ABS,
+            },
+            {
+                "category_id": 0,
+                "bbox": [2, 2, 4, 3],
+                "bbox_mode": BoxMode.XYWH_ABS,
+            },
+            {
+                "category_id": 0,
+                "bbox": [6, 5, 3, 2],
+                "bbox_mode": BoxMode.XYWH_ABS,
+            },
+        ]
+
+        img = np.ones([img_h, img_w, 3]) * 3
+
+        inputs = AugInput(image=img)
+        # pass additional arguments
+        inputs.annotations = annotations
+        transforms = AugmentationList([aug])(inputs)
+
+        self.assertIn(
+            inputs.image.shape, [torch.Size([3, 4, 3]), torch.Size([2, 3, 3])]
+        )
+
+        # from dataset mapper unused annotations will be filtered out due to the
+        #  iscrowd flag
+        image_shape = inputs.image.shape[:2]
+        annos = [
+            du.transform_instance_annotations(
+                obj,
+                transforms,
+                image_shape,
+            )
+            for obj in annotations
+            if obj.get("iscrowd", 0) == 0
+        ]
+        instances = du.annotations_to_instances(annos, image_shape)
+        filtered_instances = du.filter_empty_instances(instances)
+        self.assertEqual(len(filtered_instances), 1)
+        self.assertArrayEqual(
+            filtered_instances.gt_boxes.tensor.tolist(),
+            [[0, 0, image_shape[1], image_shape[0]]],
+        )
+
+    def assertArrayEqual(self, a1, a2):
+        self.assertTrue(np.array_equal(a1, a2))

tests/test_data_transforms_tensor.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import numpy as np
+import torch
+from d2go.data.transforms import tensor as tensor_aug
+from detectron2.data.transforms.augmentation import AugmentationList
+
+
+class TestDataTransformsTenaor(unittest.TestCase):
+    def test_tensor_aug(self):
+        """Data augmentation that that allows torch.Tensor as input"""
+
+        img = torch.ones(3, 8, 6)
+        augs = [tensor_aug.Tensor2Array(), tensor_aug.Array2Tensor()]
+
+        inputs = tensor_aug.AugInput(image=img)
+        transforms = AugmentationList(augs)(inputs)
+        self.assertArrayEqual(img, inputs.image)
+
+        # inverse is the same as itself
+        out_img = transforms.inverse().apply_image(img)
+        self.assertArrayEqual(img, out_img)
+
+    def assertArrayEqual(self, a1, a2):
+        self.assertTrue(np.array_equal(a1, a2))

tests/test_kmeans_anchors.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import os
+import unittest
+
+import numpy as np
+import torch
+from detectron2.data import DatasetCatalog, DatasetFromList, MapDataset
+from detectron2.engine import SimpleTrainer
+from d2go.modeling.kmeans_anchors import (
+    add_kmeans_anchors_cfg,
+    compute_kmeans_anchors,
+    compute_kmeans_anchors_hook,
+)
+from d2go.runner import GeneralizedRCNNRunner
+from mobile_cv.common.misc.file_utils import make_temp_directory
+from torch.utils.data.sampler import BatchSampler, Sampler
+
+from d2go.tests.data_loader_helper import LocalImageGenerator, register_toy_dataset
+
+
+class IntervalSampler(Sampler):
+    def __init__(self, size: int, interval: int):
+        self._local_indices = range(0, size, interval)
+
+    def __iter__(self):
+        yield from self._local_indices
+
+    def __len__(self):
+        return len(self._local_indices)
+
+
+def build_sequence_loader(cfg, dataset_name, mapper, total_samples, batch_size=1):
+    """
+    Similar to `build_detection_test_loader` in the way that its sampler
+        samples dataset_dicts in order and only loops once.
+    """
+    dataset_dicts = DatasetCatalog.get(dataset_name)
+    dataset = DatasetFromList(dataset_dicts)
+    dataset = MapDataset(dataset, mapper)
+
+    interval = max(1, int(len(dataset) / total_samples))
+    sampler = IntervalSampler(len(dataset), interval)
+    batch_sampler = BatchSampler(sampler, batch_size, drop_last=False)
+
+    def _trivial_batch_collator(batch):
+        return batch
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=cfg.DATALOADER.NUM_WORKERS,
+        batch_sampler=batch_sampler,
+        collate_fn=_trivial_batch_collator,
+    )
+    return data_loader
+
+
+class TestKmeansAnchors(unittest.TestCase):
+    def setUp(self):
+        self.runner = GeneralizedRCNNRunner()
+
+    def _get_default_cfg(self):
+        cfg = self.runner.get_default_cfg()
+        add_kmeans_anchors_cfg(cfg)
+        return cfg
+
+    @unittest.skip("This can only run locally and takes significant of time")
+    def test_matching_previous_results(self):
+        cfg = self._get_default_cfg()
+        cfg.INPUT.MIN_SIZE_TRAIN = (144,)
+        cfg.MODEL.KMEANS_ANCHORS.KMEANS_ANCHORS_ON = True
+        cfg.MODEL.KMEANS_ANCHORS.NUM_CLUSTERS = 10
+        cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG = 512
+        cfg.MODEL.KMEANS_ANCHORS.DATASETS = ()
+
+        # NOTE: create a data loader that samples exact the same as previous
+        # implementation. In D2Go, we will rely on the train loader instead.
+        # NOTE: in order to load OV580_XRM dataset, change the IM_DIR to:
+        # "/mnt/vol/gfsai-east/aml/mobile-vision//dataset/oculus/hand_tracking//torch/Segmentation/OV580_XRM_640x480_V3_new_rerun/images"  # noqa
+        data_loader = build_sequence_loader(
+            cfg,
+            # dataset_name="coco_2014_valminusminival",
+            # dataset_name="OV580_XRM_640x480_V3_train",
+            dataset_name="OV580_XRM_640x480_V3_heldOut_small_512",
+            mapper=self.runner.get_mapper(cfg, is_train=True),
+            total_samples=cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG,
+            batch_size=3,
+        )
+
+        kmeans_anchors = compute_kmeans_anchors(
+            cfg, data_loader, sort_by_area=False, _stride=16, _legacy_plus_one=True
+        )
+
+        # Taken from D9849940
+        reference_anchors = np.array(
+            [
+                [-15.33554182, -15.29361029, 31.33554182, 31.29361029],  # noqa
+                [-9.34156693, -9.32553548, 25.34156693, 25.32553548],  # noqa
+                [-6.03052776, -6.02034167, 22.03052776, 22.02034167],  # noqa
+                [-2.25951741, -2.182888, 18.25951741, 18.182888],  # noqa
+                [-18.93553378, -18.93553403, 34.93553378, 34.93553403],  # noqa
+                [-12.69068356, -12.73989029, 28.69068356, 28.73989029],  # noqa
+                [-24.73489189, -24.73489246, 40.73489189, 40.73489246],  # noqa
+                [-4.06014466, -4.06014469, 20.06014466, 20.06014469],  # noqa
+                [-7.61036119, -7.60467538, 23.61036119, 23.60467538],  # noqa
+                [-10.88200579, -10.87634414, 26.88200579, 26.87634414],  # noqa
+            ]
+        )
+        np.testing.assert_allclose(kmeans_anchors, reference_anchors, atol=1e-6)
+
+    def test_build_model(self):
+        cfg = self._get_default_cfg()
+        cfg.INPUT.MIN_SIZE_TRAIN = (60,)
+        cfg.MODEL.KMEANS_ANCHORS.KMEANS_ANCHORS_ON = True
+        cfg.MODEL.KMEANS_ANCHORS.NUM_CLUSTERS = 3
+        cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG = 5
+        cfg.MODEL.KMEANS_ANCHORS.DATASETS = ("toy_dataset",)
+
+        cfg.MODEL.DEVICE = "cpu"
+        cfg.MODEL.ANCHOR_GENERATOR.NAME = "KMeansAnchorGenerator"
+
+        with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir:
+            image_dir = os.path.join(dataset_dir, "images")
+            os.makedirs(image_dir)
+            image_generator = LocalImageGenerator(image_dir, width=80, height=60)
+            with register_toy_dataset(
+                "toy_dataset",
+                image_generator,
+                num_images=cfg.MODEL.KMEANS_ANCHORS.NUM_TRAINING_IMG,
+            ):
+                model = self.runner.build_model(cfg)
+                trainer = SimpleTrainer(model, data_loader=[], optimizer=None)
+                trainer_hooks = [compute_kmeans_anchors_hook(self.runner, cfg)]
+                trainer.register_hooks(trainer_hooks)
+                trainer.before_train()
+                anchor_generator = model.proposal_generator.anchor_generator
+                cell_anchors = [x for x in anchor_generator.cell_anchors]
+                gt_anchors = np.array(
+                    [
+                        [-20, -15, 20, 15]  # toy_dataset's bbox is half size of image
+                        for _ in range(cfg.MODEL.KMEANS_ANCHORS.NUM_CLUSTERS)
+                    ]
+                )
+                np.testing.assert_allclose(cell_anchors[0], gt_anchors)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_model_zoo.py

+import os
+import unittest
+
+from d2go.model_zoo import model_zoo
+
+OSSRUN = os.getenv('OSSRUN') == '1'
+
+class TestD2GoModelZoo(unittest.TestCase):
+    @unittest.skipIf(not OSSRUN, "OSS test only")
+    def test_model_zoo_pretrained(self):
+        configs = list(model_zoo._ModelZooUrls.CONFIG_PATH_TO_URL_SUFFIX.keys())
+        for cfgfile in configs:
+            model = model_zoo.get(cfgfile, trained=True)
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_modeling_image_pooler.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import d2go.data.transforms.box_utils as bu
+import d2go.modeling.image_pooler as image_pooler
+import numpy as np
+import torch
+from detectron2.structures import Boxes
+
+from d2go.tests import rcnn_helper as rh
+
+
+class TestModelingImagePooler(unittest.TestCase):
+    def test_image_pooler(self):
+        H, W = 8, 6
+        image = torch.zeros(3, H, W)
+        # xyxy
+        boxes = torch.Tensor([[2, 3, 5, 7]])
+        image[0, 3:7, 2:5] = 1
+        image[1, 3:7, 2:5] = 2
+        image[2, 3:7, 2:5] = 4
+
+        img_pooler = image_pooler.ImagePooler(resize_short=6, resize_max=12).eval()
+        pooled_img, pooled_box, transforms = img_pooler(image, boxes)
+
+        # check pooled images
+        self.assertEqual(pooled_img.shape, torch.Size([3, 8, 6]))
+        self.assertArrayEqual(torch.unique(pooled_img[0, :, :]), [1])
+        self.assertArrayEqual(torch.unique(pooled_img[1, :, :]), [2])
+        self.assertArrayEqual(torch.unique(pooled_img[2, :, :]), [4])
+
+        # check pooled boxes, in xyxy format
+        self.assertArrayEqual(pooled_box, [[0, 0, 6, 8]])
+
+        # inverse of transforms
+        trans_inv = transforms.inverse()
+
+        # inverse of boxes, xyxy
+        inversed_box = trans_inv.apply_box(pooled_box)
+        self.assertArrayEqual(inversed_box, boxes)
+
+        pooled_sub_box = np.array([[2, 2, 4, 6]])
+        inversed_sub_box = trans_inv.apply_box(pooled_sub_box)
+        self.assertArrayEqual(inversed_sub_box, [[3, 4, 4, 6]])
+
+    def test_image_pooler_scale_box(self):
+        H, W = 8, 6
+        image = torch.zeros(3, H, W)
+        # xyxy
+        boxes = torch.Tensor([[2, 3, 5, 7]])
+        image[0, 3:7, 2:5] = 1
+        image[1, 3:7, 2:5] = 2
+        image[2, 3:7, 2:5] = 4
+
+        img_pooler = image_pooler.ImagePooler(
+            resize_type=None, box_scale_factor=4.0
+        ).eval()
+        pooled_img, pooled_box, transforms = img_pooler(image, boxes)
+
+        # check pooled images
+        self.assertEqual(pooled_img.shape, torch.Size([3, 8, 6]))
+        self.assertArrayEqual(pooled_img, image)
+
+        # check pooled boxes, in xyxy format, the box before scaling
+        self.assertArrayEqual(pooled_box, [[2, 3, 5, 7]])
+
+    def test_image_pooler_scale_box_large_crop_only(self):
+        """ Crop bbox """
+        H, W = 398, 224
+        all_boxes = Boxes(torch.Tensor([[50, 40, 100, 80], [150, 60, 200, 120]]))
+        image = rh.get_batched_inputs(1, (H, W), (H, W), all_boxes)[0]["image"]
+
+        boxes = bu.get_box_union(all_boxes)
+        self.assertArrayEqual(boxes.tensor, [[50, 40, 200, 120]])
+
+        img_pooler = image_pooler.ImagePooler(
+            resize_type=None, box_scale_factor=1.0
+        ).eval()
+        pooled_img, pooled_box, transforms = img_pooler(image, boxes.tensor)
+        self.assertEqual(pooled_img.shape, torch.Size([3, 80, 150]))
+        sub_boxes = rh.get_detected_instances_from_image([{"image": pooled_img}])[
+            0
+        ].pred_boxes
+        self.assertArrayEqual(sub_boxes.tensor, [[0, 0, 50, 40], [100, 20, 150, 80]])
+
+    def test_image_pooler_scale_box_large_crop_and_scale(self):
+        """ Crop bbox that is scaled """
+        H, W = 398, 224
+        all_boxes = Boxes(torch.Tensor([[50, 40, 100, 80], [150, 60, 200, 120]]))
+        image = rh.get_batched_inputs(1, (H, W), (H, W), all_boxes)[0]["image"]
+        boxes = bu.get_box_union(all_boxes)
+
+        img_pooler = image_pooler.ImagePooler(
+            resize_type=None, box_scale_factor=1.2
+        ).eval()
+        pooled_img, pooled_box, transforms = img_pooler(image, boxes.tensor)
+        self.assertEqual(pooled_img.shape, torch.Size([3, 96, 180]))
+
+        # bbox with scaling in the original space
+        orig_crop_box = transforms.inverse().apply_box(
+            [0, 0, pooled_img.shape[2], pooled_img.shape[1]]
+        )
+        self.assertArrayEqual(orig_crop_box, [[35, 32, 215, 128]])
+
+        sub_boxes = rh.get_detected_instances_from_image([{"image": pooled_img}])[
+            0
+        ].pred_boxes
+        # gt_offset_xy = (50 - 35 = 15, 40 - 32 = 8)
+        self.assertArrayEqual(sub_boxes.tensor, [[15, 8, 65, 48], [115, 28, 165, 88]])
+
+    def test_image_pooler_scale_box_large_crop_scale_and_resize(self):
+        """ Crop bbox that is scaled, resize the cropped box """
+        H, W = 398, 224
+        all_boxes = Boxes(torch.Tensor([[50, 40, 100, 80], [150, 60, 200, 120]]))
+        image = rh.get_batched_inputs(1, (H, W), (H, W), all_boxes)[0]["image"]
+        boxes = bu.get_box_union(all_boxes)
+
+        img_pooler = image_pooler.ImagePooler(
+            resize_type="resize_shortest",
+            resize_short=48,
+            resize_max=180,
+            box_scale_factor=1.2,
+        ).eval()
+        pooled_img, pooled_box, transforms = img_pooler(image, boxes.tensor)
+        self.assertEqual(pooled_img.shape, torch.Size([3, 48, 90]))
+
+        # bbox with scaling in the original space
+        orig_crop_box = transforms.inverse().apply_box(
+            [0, 0, pooled_img.shape[2], pooled_img.shape[1]]
+        )
+        self.assertArrayEqual(orig_crop_box, [[35, 32, 215, 128]])
+
+        # bbox without scaling in the original space
+        orig_boxes = transforms.inverse().apply_box(pooled_box)
+        self.assertArrayEqual(orig_boxes, boxes.tensor)
+
+        sub_boxes = rh.get_detected_instances_from_image([{"image": pooled_img}])[
+            0
+        ].pred_boxes
+        # [[7.5, 4, 32.5, 24], [57.5, 14, 82.5, 44]]
+        self.assertArrayEqual(sub_boxes.tensor, [[7, 4, 33, 24], [57, 14, 83, 44]])
+
+    def assertArrayEqual(self, a1, a2):
+        self.assertTrue(np.array_equal(a1, a2))

tests/test_modeling_model_ema.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import copy
+import itertools
+import unittest
+
+import torch
+from d2go.modeling import model_ema
+import d2go.runner.default_runner as default_runner
+
+from d2go.tests import helper
+
+
+class TestArch(torch.nn.Module):
+    def __init__(self, value=None, int_value=None):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 4, kernel_size=3, stride=1, padding=1)
+        self.bn = torch.nn.BatchNorm2d(4)
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))
+        if value is not None:
+            self.set_const_weights(value, int_value)
+
+    def forward(self, x):
+        ret = self.conv(x)
+        ret = self.bn(ret)
+        ret = self.relu(ret)
+        ret = self.avgpool(ret)
+        return ret
+
+    def set_const_weights(self, value, int_value=None):
+        if int_value is None:
+            int_value = int(value)
+        for x in itertools.chain(self.parameters(), self.buffers()):
+            if x.dtype == torch.float32:
+                x.data.fill_(value)
+            else:
+                x.data.fill_(int_value)
+
+
+def _compare_state_dict(model1, model2, abs_error=1e-3):
+    sd1 = model1.state_dict()
+    sd2 = model2.state_dict()
+    if len(sd1) != len(sd2):
+        return False
+    if set(sd1.keys()) != set(sd2.keys()):
+        return False
+    for name in sd1:
+        if sd1[name].dtype == torch.float32:
+            if torch.abs((sd1[name] - sd2[name])).max() > abs_error:
+                return False
+        elif (sd1[name] != sd2[name]).any():
+            return False
+    return True
+
+
+class TestModelingModelEMA(unittest.TestCase):
+    def test_emastate(self):
+        model = TestArch()
+        state = model_ema.EMAState.FromModel(model)
+        # two for conv (conv.weight, conv.bias),
+        # five for bn (bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.num_batches_tracked)
+        self.assertEqual(len(state.state), 7)
+
+        for _, val in state.state.items():
+            self.assertFalse(val.requires_grad)
+
+        model1 = TestArch()
+        self.assertFalse(_compare_state_dict(model, model1))
+
+        state.apply_to(model1)
+        self.assertTrue(_compare_state_dict(model, model1))
+
+    def test_emastate_saveload(self):
+        model = TestArch()
+        state = model_ema.EMAState.FromModel(model)
+
+        model1 = TestArch()
+        self.assertFalse(_compare_state_dict(model, model1))
+
+        state1 = model_ema.EMAState()
+        state1.load_state_dict(state.state_dict())
+        state1.apply_to(model1)
+        self.assertTrue(_compare_state_dict(model, model1))
+
+    @helper.skip_if_no_gpu
+    def test_emastate_crossdevice(self):
+        model = TestArch()
+        model.cuda()
+        # state on gpu
+        state = model_ema.EMAState.FromModel(model)
+        self.assertEqual(state.device, torch.device("cuda:0"))
+        # target model on cpu
+        model1 = TestArch()
+        state.apply_to(model1)
+        self.assertEqual(next(model1.parameters()).device, torch.device("cpu"))
+        self.assertTrue(_compare_state_dict(copy.deepcopy(model).cpu(), model1))
+
+        # state on cpu
+        state1 = model_ema.EMAState.FromModel(model, device="cpu")
+        self.assertEqual(state1.device, torch.device("cpu"))
+        # target model on gpu
+        model2 = TestArch()
+        model2.cuda()
+        state1.apply_to(model2)
+        self.assertEqual(next(model2.parameters()).device, torch.device("cuda:0"))
+        self.assertTrue(_compare_state_dict(model, model2))
+
+    def test_ema_updater(self):
+        model = TestArch()
+        state = model_ema.EMAState()
+
+        updated_model = TestArch()
+
+        updater = model_ema.EMAUpdater(state, decay=0.0)
+        updater.init_state(model)
+        for _ in range(3):
+            cur = TestArch()
+            updater.update(cur)
+            state.apply_to(updated_model)
+            # weight decay == 0.0, always use new model
+            self.assertTrue(_compare_state_dict(updated_model, cur))
+
+        updater = model_ema.EMAUpdater(state, decay=1.0)
+        updater.init_state(model)
+        for _ in range(3):
+            cur = TestArch()
+            updater.update(cur)
+            state.apply_to(updated_model)
+            # weight decay == 1.0, always use init model
+            self.assertTrue(_compare_state_dict(updated_model, model))
+
+    def test_ema_updater_decay(self):
+        state = model_ema.EMAState()
+
+        updater = model_ema.EMAUpdater(state, decay=0.7)
+        updater.init_state(TestArch(1.0))
+        gt_val = 1.0
+        gt_val_int = 1
+        for idx in range(3):
+            updater.update(TestArch(float(idx)))
+            updated_model = state.get_ema_model(TestArch())
+            gt_val = gt_val * 0.7 + float(idx) * 0.3
+            gt_val_int = int(gt_val_int * 0.7 + float(idx) * 0.3)
+            self.assertTrue(
+                _compare_state_dict(updated_model, TestArch(gt_val, gt_val_int))
+            )
+
+
+class TestModelingModelEMAHook(unittest.TestCase):
+    def test_ema_hook(self):
+        runner = default_runner.Detectron2GoRunner()
+        cfg = runner.get_default_cfg()
+        cfg.MODEL.DEVICE = "cpu"
+        cfg.MODEL_EMA.ENABLED = True
+        # use new model weights
+        cfg.MODEL_EMA.DECAY = 0.0
+
+        model = TestArch()
+        model_ema.may_build_model_ema(cfg, model)
+        self.assertTrue(hasattr(model, "ema_state"))
+
+        ema_hook = model_ema.EMAHook(cfg, model)
+        ema_hook.before_train()
+        ema_hook.before_step()
+        model.set_const_weights(2.0)
+        ema_hook.after_step()
+        ema_hook.after_train()
+
+        ema_checkpointers = model_ema.may_get_ema_checkpointer(cfg, model)
+        self.assertEqual(len(ema_checkpointers), 1)
+
+        out_model = TestArch()
+        ema_checkpointers["ema_state"].apply_to(out_model)
+        self.assertTrue(
+            _compare_state_dict(out_model, model)
+        )

tests/test_nms.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import unittest
+
+import numpy as np
+import torch
+from detectron2.layers import nms as box_nms
+
+
+class TestNMS(unittest.TestCase):
+    def test_nms_cpu(self):
+        """ Match unit test UtilsNMSTest.TestNMS in
+            caffe2/operators/generate_proposals_op_util_nms_test.cc
+        """
+
+        inputs = (
+            np.array(
+                [
+                    10,
+                    10,
+                    50,
+                    60,
+                    0.5,
+                    11,
+                    12,
+                    48,
+                    60,
+                    0.7,
+                    8,
+                    9,
+                    40,
+                    50,
+                    0.6,
+                    100,
+                    100,
+                    150,
+                    140,
+                    0.9,
+                    99,
+                    110,
+                    155,
+                    139,
+                    0.8,
+                ]
+            )
+            .astype(np.float32)
+            .reshape(-1, 5)
+        )
+
+        boxes = torch.from_numpy(inputs[:, :4])
+        scores = torch.from_numpy(inputs[:, 4])
+        test_thresh = [0.1, 0.3, 0.5, 0.8, 0.9]
+        gt_indices = [[1, 3], [1, 3], [1, 3], [1, 2, 3, 4], [0, 1, 2, 3, 4]]
+
+        for thresh, gt_index in zip(test_thresh, gt_indices):
+            keep_indices = box_nms(boxes, scores, thresh)
+            keep_indices = np.sort(keep_indices)
+            np.testing.assert_array_equal(keep_indices, np.array(gt_index))
+
+    def test_nms1_cpu(self):
+        """ Match unit test UtilsNMSTest.TestNMS1 in
+            caffe2/operators/generate_proposals_op_util_nms_test.cc
+        """
+
+        boxes = torch.from_numpy(
+            np.array(
+                [
+                    [350.9821, 161.8200, 369.9685, 205.2372],
+                    [250.5236, 154.2844, 274.1773, 204.9810],
+                    [471.4920, 160.4118, 496.0094, 213.4244],
+                    [352.0421, 164.5933, 366.4458, 205.9624],
+                    [166.0765, 169.7707, 183.0102, 232.6606],
+                    [252.3000, 183.1449, 269.6541, 210.6747],
+                    [469.7862, 162.0192, 482.1673, 187.0053],
+                    [168.4862, 174.2567, 181.7437, 232.9379],
+                    [470.3290, 162.3442, 496.4272, 214.6296],
+                    [251.0450, 155.5911, 272.2693, 203.3675],
+                    [252.0326, 154.7950, 273.7404, 195.3671],
+                    [351.7479, 161.9567, 370.6432, 204.3047],
+                    [496.3306, 161.7157, 515.0573, 210.7200],
+                    [471.0749, 162.6143, 485.3374, 207.3448],
+                    [250.9745, 160.7633, 264.1924, 206.8350],
+                    [470.4792, 169.0351, 487.1934, 220.2984],
+                    [474.4227, 161.9546, 513.1018, 215.5193],
+                    [251.9428, 184.1950, 262.6937, 207.6416],
+                    [252.6623, 175.0252, 269.8806, 213.7584],
+                    [260.9884, 157.0351, 288.3554, 206.6027],
+                    [251.3629, 164.5101, 263.2179, 202.4203],
+                    [471.8361, 190.8142, 485.6812, 220.8586],
+                    [248.6243, 156.9628, 264.3355, 199.2767],
+                    [495.1643, 158.0483, 512.6261, 184.4192],
+                    [376.8718, 168.0144, 387.3584, 201.3210],
+                    [122.9191, 160.7433, 172.5612, 231.3837],
+                    [350.3857, 175.8806, 366.2500, 205.4329],
+                    [115.2958, 162.7822, 161.9776, 229.6147],
+                    [168.4375, 177.4041, 180.8028, 232.4551],
+                    [169.7939, 184.4330, 181.4767, 232.1220],
+                    [347.7536, 175.9356, 355.8637, 197.5586],
+                    [495.5434, 164.6059, 516.4031, 207.7053],
+                    [172.1216, 194.6033, 183.1217, 235.2653],
+                    [264.2654, 181.5540, 288.4626, 214.0170],
+                    [111.7971, 183.7748, 137.3745, 225.9724],
+                    [253.4919, 186.3945, 280.8694, 210.0731],
+                    [165.5334, 169.7344, 185.9159, 232.8514],
+                    [348.3662, 184.5187, 354.9081, 201.4038],
+                    [164.6562, 162.5724, 186.3108, 233.5010],
+                    [113.2999, 186.8410, 135.8841, 219.7642],
+                    [117.0282, 179.8009, 142.5375, 221.0736],
+                    [462.1312, 161.1004, 495.3576, 217.2208],
+                    [462.5800, 159.9310, 501.2937, 224.1655],
+                    [503.5242, 170.0733, 518.3792, 209.0113],
+                    [250.3658, 195.5925, 260.6523, 212.4679],
+                    [108.8287, 163.6994, 146.3642, 229.7261],
+                    [256.7617, 187.3123, 288.8407, 211.2013],
+                    [161.2781, 167.4801, 186.3751, 232.7133],
+                    [115.3760, 177.5859, 163.3512, 236.9660],
+                    [248.9077, 188.0919, 264.8579, 207.9718],
+                    [108.1349, 160.7851, 143.6370, 229.6243],
+                    [465.0900, 156.7555, 490.3561, 213.5704],
+                    [107.5338, 173.4323, 141.0704, 235.2910],
+                ]
+            ).astype(np.float32)
+        )
+        scores = torch.from_numpy(
+            np.array(
+                [
+                    0.1919,
+                    0.3293,
+                    0.0860,
+                    0.1600,
+                    0.1885,
+                    0.4297,
+                    0.0974,
+                    0.2711,
+                    0.1483,
+                    0.1173,
+                    0.1034,
+                    0.2915,
+                    0.1993,
+                    0.0677,
+                    0.3217,
+                    0.0966,
+                    0.0526,
+                    0.5675,
+                    0.3130,
+                    0.1592,
+                    0.1353,
+                    0.0634,
+                    0.1557,
+                    0.1512,
+                    0.0699,
+                    0.0545,
+                    0.2692,
+                    0.1143,
+                    0.0572,
+                    0.1990,
+                    0.0558,
+                    0.1500,
+                    0.2214,
+                    0.1878,
+                    0.2501,
+                    0.1343,
+                    0.0809,
+                    0.1266,
+                    0.0743,
+                    0.0896,
+                    0.0781,
+                    0.0983,
+                    0.0557,
+                    0.0623,
+                    0.5808,
+                    0.3090,
+                    0.1050,
+                    0.0524,
+                    0.0513,
+                    0.4501,
+                    0.4167,
+                    0.0623,
+                    0.1749,
+                ]
+            ).astype(np.float32)
+        )
+
+        gt_indices = np.array(
+            [
+                1,
+                6,
+                7,
+                8,
+                11,
+                12,
+                13,
+                14,
+                17,
+                18,
+                19,
+                21,
+                23,
+                24,
+                25,
+                26,
+                30,
+                32,
+                33,
+                34,
+                35,
+                37,
+                43,
+                44,
+                47,
+                50,
+            ]
+        )
+        keep_indices = box_nms(boxes, scores, 0.5)
+        keep_indices = np.sort(keep_indices)
+
+        np.testing.assert_array_equal(keep_indices, gt_indices)
+
+
+if __name__ == "__main__":
+    unittest.main()