first commit

ssd/config/__init__.py

@@ -0,0 +1 @@
+from .defaults import _C as cfg

ssd/config/defaults.py

@@ -0,0 +1,97 @@
+from yacs.config import CfgNode as CN
+
+_C = CN()
+
+_C.MODEL = CN()
+_C.MODEL.META_ARCHITECTURE = 'SSDDetector'
+_C.MODEL.DEVICE = "cuda"
+# match default boxes to any ground truth with jaccard overlap higher than a threshold (0.5)
+_C.MODEL.THRESHOLD = 0.5
+_C.MODEL.NUM_CLASSES = 21
+# Hard negative mining
+_C.MODEL.NEG_POS_RATIO = 3
+_C.MODEL.CENTER_VARIANCE = 0.1
+_C.MODEL.SIZE_VARIANCE = 0.2
+
+# ---------------------------------------------------------------------------- #
+# Backbone
+# ---------------------------------------------------------------------------- #
+_C.MODEL.BACKBONE = CN()
+_C.MODEL.BACKBONE.NAME = 'vgg'
+_C.MODEL.BACKBONE.OUT_CHANNELS = (512, 1024, 512, 256, 256, 256)
+_C.MODEL.BACKBONE.PRETRAINED = True
+
+# -----------------------------------------------------------------------------
+# PRIORS
+# -----------------------------------------------------------------------------
+_C.MODEL.PRIORS = CN()
+_C.MODEL.PRIORS.FEATURE_MAPS = [38, 19, 10, 5, 3, 1]
+_C.MODEL.PRIORS.STRIDES = [8, 16, 32, 64, 100, 300]
+_C.MODEL.PRIORS.MIN_SIZES = [30, 60, 111, 162, 213, 264]
+_C.MODEL.PRIORS.MAX_SIZES = [60, 111, 162, 213, 264, 315]
+_C.MODEL.PRIORS.ASPECT_RATIOS = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
+# When has 1 aspect ratio, every location has 4 boxes, 2 ratio 6 boxes.
+# #boxes = 2 + #ratio * 2
+_C.MODEL.PRIORS.BOXES_PER_LOCATION = [4, 6, 6, 6, 4, 4]  # number of boxes per feature map location
+_C.MODEL.PRIORS.CLIP = True
+
+# -----------------------------------------------------------------------------
+# Box Head
+# -----------------------------------------------------------------------------
+_C.MODEL.BOX_HEAD = CN()
+_C.MODEL.BOX_HEAD.NAME = 'SSDBoxHead'
+_C.MODEL.BOX_HEAD.PREDICTOR = 'SSDBoxPredictor'
+
+# -----------------------------------------------------------------------------
+# INPUT
+# -----------------------------------------------------------------------------
+_C.INPUT = CN()
+# Image size
+_C.INPUT.IMAGE_SIZE = 300
+# Values to be used for image normalization, RGB layout
+_C.INPUT.PIXEL_MEAN = [123, 117, 104]
+
+# -----------------------------------------------------------------------------
+# Dataset
+# -----------------------------------------------------------------------------
+_C.DATASETS = CN()
+# List of the dataset names for training, as present in paths_catalog.py
+_C.DATASETS.TRAIN = ()
+# List of the dataset names for testing, as present in paths_catalog.py
+_C.DATASETS.TEST = ()
+_C.DATASETS.INF = ()
+
+# -----------------------------------------------------------------------------
+# DataLoader
+# -----------------------------------------------------------------------------
+_C.DATA_LOADER = CN()
+# Number of data loading threads
+_C.DATA_LOADER.NUM_WORKERS = 8
+_C.DATA_LOADER.PIN_MEMORY = True
+
+# ---------------------------------------------------------------------------- #
+# Solver
+# ---------------------------------------------------------------------------- #
+_C.SOLVER = CN()
+# train configs
+_C.SOLVER.MAX_ITER = 120000
+_C.SOLVER.LR_STEPS = [80000, 100000]
+_C.SOLVER.GAMMA = 0.1
+_C.SOLVER.BATCH_SIZE = 4#32
+_C.SOLVER.LR = 1e-3
+_C.SOLVER.MOMENTUM = 0.9
+_C.SOLVER.WEIGHT_DECAY = 5e-4
+_C.SOLVER.WARMUP_FACTOR = 1.0 / 3
+_C.SOLVER.WARMUP_ITERS = 500
+
+# ---------------------------------------------------------------------------- #
+# Specific test options
+# ---------------------------------------------------------------------------- #
+_C.TEST = CN()
+_C.TEST.NMS_THRESHOLD = 0.4 #0.45
+_C.TEST.CONFIDENCE_THRESHOLD = 0.3#0.01
+_C.TEST.MAX_PER_CLASS = -1
+_C.TEST.MAX_PER_IMAGE = 100
+_C.TEST.BATCH_SIZE = 1
+
+_C.OUTPUT_DIR = 'outputs'

ssd/config/hat-512.yaml

@@ -0,0 +1,30 @@
+MODEL:
+  NUM_CLASSES: 21
+  BACKBONE:
+    OUT_CHANNELS: (512, 1024, 512, 256, 256, 256, 256)
+  PRIORS:
+    FEATURE_MAPS: [64, 32, 16, 8, 4, 2, 1]
+    STRIDES: [8, 16, 32, 64, 128, 256, 512]
+    MIN_SIZES: [35.84, 76.8, 153.6, 230.4, 307.2, 384.0, 460.8]
+    MAX_SIZES: [76.8, 153.6, 230.4, 307.2, 384.0, 460.8, 537.65]
+    ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]]
+    BOXES_PER_LOCATION: [4, 6, 6, 6, 6, 4, 4]
+INPUT:
+  IMAGE_SIZE: 512
+DATASETS:
+  #TRAIN: ("voc_2007_birdtrain", "voc_2007_birdval")
+  #TRAIN: ("voc_2007_trainval", )
+  #TEST: ("voc_2007_test", )
+  #TEST: ("vocbird_2019_test", )
+  TRAIN: ("voc_3_train",)
+  TEST: ("voc_3_test", )
+  INF: ("inference_hatcher",)
+SOLVER:
+  MAX_ITER: 175000
+
+  LR_STEPS: [80000, 100000]
+  GAMMA: 0.1
+  BATCH_SIZE: 16
+  LR: 1e-5
+#1e-3
+OUTPUT_DIR: '/data1/Hatcher/model/data_MPKFNet/weigth/detection/vgg_ssd512_voc0712VOC3/'

ssd/config/path_catlog.py

@@ -0,0 +1,114 @@
+import os
+
+
+class DatasetCatalog:
+    DATA_DIR = ''
+    DATASETS = {
+        'voc_ICCV_train': {
+            "data_dir": "VOCiccvtrain",
+            "split": "trainyolo"
+        },
+        'voc_ICCV_test': {
+            "data_dir": "VOCiccvtest",
+            "split": "testyolo"
+        },
+        'voc_2007_train': {
+            "data_dir": "VOC2007",
+            "split": "train"
+        },
+        'voc_3_train': {
+            "data_dir": "VOC3",
+            "split": "train"
+        },
+        'voc_3_test': {
+            "data_dir": "VOC3",
+            "split": "test"
+        },
+        'voc_2007_val': {
+            "data_dir": "VOC2007",
+            "split": "val"
+        },
+        'voc_2007_trainval': {
+            "data_dir": "VOC2007",
+            "split": "trainval"
+        },
+        'voc_2007_test': {
+            "data_dir": "VOC2007",
+            "split": "test"
+        },
+        'vocbird_2019_test': {
+            "data_dir": "HatcherW",
+            "split": "test"
+        },
+        'voc_2012_train': {
+            "data_dir": "VOC2012",
+            "split": "train"
+        },
+        'voc_2012_val': {
+            "data_dir": "VOC2012",
+            "split": "val"
+        },
+        'voc_2012_trainval': {
+            "data_dir": "VOC2012",
+            "split": "trainval"
+        },
+        'voc_2012_test': {
+            "data_dir": "VOC2012",
+            "split": "test"
+        },
+        'coco_2014_valminusminival': {
+            "data_dir": "val2014",
+            "ann_file": "annotations/instances_valminusminival2014.json"
+        },
+        'coco_2014_minival': {
+            "data_dir": "val2014",
+            "ann_file": "annotations/instances_minival2014.json"
+        },
+        'coco_2014_train': {
+            "data_dir": "train2014",
+            "ann_file": "annotations/instances_train2014.json"
+        },
+        'coco_2014_val': {
+            "data_dir": "val2014",
+            "ann_file": "annotations/instances_val2014.json"
+        },
+        "voc_2007_birdtest": {
+            "data_dir": "VOC2007",
+            "split": "bird_test"
+        },
+        "voc_2007_birdval": {
+            "data_dir": "VOC2007",
+            "split": "bird_trainval"
+        },
+        "voc_2007_birdtrain": {
+            "data_dir": "VOC2007",
+            "split": "bird_train"
+        },
+    }
+
+    @staticmethod
+    def get(name):
+        if "voc" in name:
+            voc_root = DatasetCatalog.DATA_DIR
+            if 'VOC_ROOT' in os.environ:
+                voc_root = os.environ['VOC_ROOT']
+
+            attrs = DatasetCatalog.DATASETS[name]
+            args = dict(
+                data_dir=os.path.join(voc_root, attrs["data_dir"]),
+                split=attrs["split"],
+            )
+            return dict(factory="VOCDataset", args=args)
+        elif "coco" in name:
+            coco_root = DatasetCatalog.DATA_DIR
+            if 'COCO_ROOT' in os.environ:
+                coco_root = os.environ['COCO_ROOT']
+
+            attrs = DatasetCatalog.DATASETS[name]
+            args = dict(
+                data_dir=os.path.join(coco_root, attrs["data_dir"]),
+                ann_file=os.path.join(coco_root, attrs["ann_file"]),
+            )
+            return dict(factory="COCODataset", args=args)
+
+        raise RuntimeError("Dataset not available: {}".format(name))

ssd/data/build.py

@@ -0,0 +1,61 @@
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import default_collate
+
+from ssd.data import samplers
+from ssd.data.datasets import build_dataset
+from ssd.data.transforms import build_transforms, build_target_transform
+from ssd.structures.container import Container
+
+
+class BatchCollator:
+    def __init__(self, is_train=True):
+        self.is_train = is_train
+
+    def __call__(self, batch):
+        transposed_batch = list(zip(*batch))
+        images = default_collate(transposed_batch[0])
+        img_ids = default_collate(transposed_batch[2])
+
+        if self.is_train:
+            list_targets = transposed_batch[1]
+            targets = Container(
+                {key: default_collate([d[key] for d in list_targets]) for key in list_targets[0]}
+            )
+        else:
+            targets = None
+        return images, targets, img_ids
+
+
+def make_data_loader(cfg, is_train=True, distributed=False, max_iter=None, start_iter=0):
+    train_transform = build_transforms(cfg, is_train=is_train)
+    target_transform = build_target_transform(cfg) if is_train else None
+    dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST
+    datasets = build_dataset(dataset_list, transform=train_transform, target_transform=target_transform, is_train=is_train)
+
+    shuffle = is_train or distributed
+
+    data_loaders = []
+
+    for dataset in datasets:
+        if distributed:
+            sampler = samplers.DistributedSampler(dataset, shuffle=shuffle)
+        elif shuffle:
+            sampler = torch.utils.data.RandomSampler(dataset)
+        else:
+            sampler = torch.utils.data.sampler.SequentialSampler(dataset)
+
+        batch_size = cfg.SOLVER.BATCH_SIZE if is_train else cfg.TEST.BATCH_SIZE
+        batch_sampler = torch.utils.data.sampler.BatchSampler(sampler=sampler, batch_size=batch_size, drop_last=False)
+        if max_iter is not None:
+            batch_sampler = samplers.IterationBasedBatchSampler(batch_sampler, num_iterations=max_iter, start_iter=start_iter)
+
+        data_loader = DataLoader(dataset, num_workers=cfg.DATA_LOADER.NUM_WORKERS, batch_sampler=batch_sampler,
+                                 pin_memory=cfg.DATA_LOADER.PIN_MEMORY, collate_fn=BatchCollator(is_train))
+        data_loaders.append(data_loader)
+
+    if is_train:
+        # during training, a single (possibly concatenated) data_loader is returned
+        assert len(data_loaders) == 1
+        return data_loaders[0]
+    return data_loaders

ssd/data/datasets/__init__.py

@@ -0,0 +1,35 @@
+from torch.utils.data import ConcatDataset
+
+from ssd.config.path_catlog import DatasetCatalog
+from .voc import VOCDataset
+from .coco import COCODataset
+
+_DATASETS = {
+    'VOCDataset': VOCDataset,
+    'COCODataset': COCODataset,
+}
+
+
+def build_dataset(dataset_list, transform=None, target_transform=None, is_train=True):
+    assert len(dataset_list) > 0
+    datasets = []
+    for dataset_name in dataset_list:
+        data = DatasetCatalog.get(dataset_name)
+        args = data['args']
+        factory = _DATASETS[data['factory']]
+        args['transform'] = transform
+        args['target_transform'] = target_transform
+        if factory == VOCDataset:
+            args['keep_difficult'] = not is_train
+        elif factory == COCODataset:
+            args['remove_empty'] = is_train
+        dataset = factory(**args)
+        datasets.append(dataset)
+    # for testing, return a list of datasets
+    if not is_train:
+        return datasets
+    dataset = datasets[0]
+    if len(datasets) > 1:
+        dataset = ConcatDataset(datasets)
+
+    return [dataset]

ssd/data/datasets/coco.py

@@ -0,0 +1,92 @@
+import os
+import torch.utils.data
+import numpy as np
+from PIL import Image
+
+from ssd.structures.container import Container
+
+
+class COCODataset(torch.utils.data.Dataset):
+    class_names = ('__background__',
+                   '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')
+
+    def __init__(self, data_dir, ann_file, transform=None, target_transform=None, remove_empty=False):
+        from pycocotools.coco import COCO
+        self.coco = COCO(ann_file)
+        self.data_dir = data_dir
+        self.transform = transform
+        self.target_transform = target_transform
+        self.remove_empty = remove_empty
+        if self.remove_empty:
+            # when training, images without annotations are removed.
+            self.ids = list(self.coco.imgToAnns.keys())
+        else:
+            # when testing, all images used.
+            self.ids = list(self.coco.imgs.keys())
+        coco_categories = sorted(self.coco.getCatIds())
+        self.coco_id_to_contiguous_id = {coco_id: i + 1 for i, coco_id in enumerate(coco_categories)}
+        self.contiguous_id_to_coco_id = {v: k for k, v in self.coco_id_to_contiguous_id.items()}
+
+    def __getitem__(self, index):
+        image_id = self.ids[index]
+        boxes, labels = self._get_annotation(image_id)
+        image = self._read_image(image_id)
+        if self.transform:
+            image, boxes, labels = self.transform(image, boxes, labels)
+        if self.target_transform:
+            boxes, labels = self.target_transform(boxes, labels)
+        targets = Container(
+            boxes=boxes,
+            labels=labels,
+        )
+        return image, targets, index
+
+    def get_annotation(self, index):
+        image_id = self.ids[index]
+        return image_id, self._get_annotation(image_id)
+
+    def __len__(self):
+        return len(self.ids)
+
+    def _get_annotation(self, image_id):
+        ann_ids = self.coco.getAnnIds(imgIds=image_id)
+        ann = self.coco.loadAnns(ann_ids)
+        # filter crowd annotations
+        ann = [obj for obj in ann if obj["iscrowd"] == 0]
+        boxes = np.array([self._xywh2xyxy(obj["bbox"]) for obj in ann], np.float32).reshape((-1, 4))
+        labels = np.array([self.coco_id_to_contiguous_id[obj["category_id"]] for obj in ann], np.int64).reshape((-1,))
+        # remove invalid boxes
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        labels = labels[keep]
+        return boxes, labels
+
+    def _xywh2xyxy(self, box):
+        x1, y1, w, h = box
+        return [x1, y1, x1 + w, y1 + h]
+
+    def get_img_info(self, index):
+        image_id = self.ids[index]
+        img_data = self.coco.imgs[image_id]
+        return img_data
+
+    def _read_image(self, image_id):
+        file_name = self.coco.loadImgs(image_id)[0]['file_name']
+        image_file = os.path.join(self.data_dir, file_name)
+        image = Image.open(image_file).convert("RGB")
+        image = np.array(image)
+        return image

ssd/data/datasets/evaluation/__init__.py

@@ -0,0 +1,24 @@
+from ssd.data.datasets import VOCDataset, COCODataset
+from .coco import coco_evaluation
+from .voc import voc_evaluation
+
+
+def evaluate(dataset, predictions, output_dir, **kwargs):
+    """evaluate dataset using different methods based on dataset type.
+    Args:
+        dataset: Dataset object
+        predictions(list[(boxes, labels, scores)]): Each item in the list represents the
+            prediction results for one image. And the index should match the dataset index.
+        output_dir: output folder, to save evaluation files or results.
+    Returns:
+        evaluation result
+    """
+    args = dict(
+        dataset=dataset, predictions=predictions, output_dir=output_dir, **kwargs,
+    )
+    if isinstance(dataset, VOCDataset):
+        return voc_evaluation(**args)
+    elif isinstance(dataset, COCODataset):
+        return coco_evaluation(**args)
+    else:
+        raise NotImplementedError

ssd/data/datasets/evaluation/coco/__init__.py

@@ -0,0 +1,62 @@
+import json
+import logging
+import os
+from datetime import datetime
+
+
+def coco_evaluation(dataset, predictions, output_dir, iteration=None):
+    coco_results = []
+    for i, prediction in enumerate(predictions):
+        img_info = dataset.get_img_info(i)
+        prediction = prediction.resize((img_info['width'], img_info['height'])).numpy()
+        boxes, labels, scores = prediction['boxes'], prediction['labels'], prediction['scores']
+
+        image_id, annotation = dataset.get_annotation(i)
+        class_mapper = dataset.contiguous_id_to_coco_id
+        if labels.shape[0] == 0:
+            continue
+
+        boxes = boxes.tolist()
+        labels = labels.tolist()
+        scores = scores.tolist()
+        coco_results.extend(
+            [
+                {
+                    "image_id": image_id,
+                    "category_id": class_mapper[labels[k]],
+                    "bbox": [box[0], box[1], box[2] - box[0], box[3] - box[1]],  # to xywh format
+                    "score": scores[k],
+                }
+                for k, box in enumerate(boxes)
+            ]
+        )
+    iou_type = 'bbox'
+    json_result_file = os.path.join(output_dir, iou_type + ".json")
+    logger = logging.getLogger("SSD.inference")
+    logger.info('Writing results to {}...'.format(json_result_file))
+    with open(json_result_file, "w") as f:
+        json.dump(coco_results, f)
+    from pycocotools.cocoeval import COCOeval
+    coco_gt = dataset.coco
+    coco_dt = coco_gt.loadRes(json_result_file)
+    coco_eval = COCOeval(coco_gt, coco_dt, iou_type)
+    coco_eval.evaluate()
+    coco_eval.accumulate()
+    coco_eval.summarize()
+
+    result_strings = []
+    keys = ["AP", "AP50", "AP75", "APs", "APm", "APl"]
+    metrics = {}
+    for i, key in enumerate(keys):
+        metrics[key] = coco_eval.stats[i]
+        logger.info('{:<10}: {}'.format(key, round(coco_eval.stats[i], 3)))
+        result_strings.append('{:<10}: {}'.format(key, round(coco_eval.stats[i], 3)))
+
+    if iteration is not None:
+        result_path = os.path.join(output_dir, 'result_{:07d}.txt'.format(iteration))
+    else:
+        result_path = os.path.join(output_dir, 'result_{}.txt'.format(datetime.now().strftime('%Y-%m-%d_%H-%M-%S')))
+    with open(result_path, "w") as f:
+        f.write('\n'.join(result_strings))
+
+    return dict(metrics=metrics)

ssd/data/datasets/evaluation/voc/__init__.py

@@ -0,0 +1,60 @@
+import logging
+import os
+from datetime import datetime
+
+import numpy as np
+
+from .eval_detection_voc import eval_detection_voc
+
+
+def voc_evaluation(dataset, predictions, output_dir, iteration=None):
+    class_names = dataset.class_names
+
+    pred_boxes_list = []
+    pred_labels_list = []
+    pred_scores_list = []
+    gt_boxes_list = []
+    gt_labels_list = []
+    gt_difficults = []
+
+    for i in range(len(dataset)):
+        image_id, annotation = dataset.get_annotation(i)
+        gt_boxes, gt_labels, is_difficult = annotation
+        gt_boxes_list.append(gt_boxes)
+        gt_labels_list.append(gt_labels)
+        gt_difficults.append(is_difficult.astype(np.bool))
+
+        img_info = dataset.get_img_info(i)
+        prediction = predictions[i]
+        prediction = prediction.resize((img_info['width'], img_info['height'])).numpy()
+        boxes, labels, scores = prediction['boxes'], prediction['labels'], prediction['scores']
+
+        pred_boxes_list.append(boxes)
+        pred_labels_list.append(labels)
+        pred_scores_list.append(scores)
+    result = eval_detection_voc(pred_bboxes=pred_boxes_list,
+                                pred_labels=pred_labels_list,
+                                pred_scores=pred_scores_list,
+                                gt_bboxes=gt_boxes_list,
+                                gt_labels=gt_labels_list,
+                                gt_difficults=gt_difficults,
+                                iou_thresh=0.5,
+                                use_07_metric=True)
+    logger = logging.getLogger("SSD.inference")
+    result_str = "mAP: {:.4f}\n".format(result["map"])
+    metrics = {'mAP': result["map"]}
+    for i, ap in enumerate(result["ap"]):
+        if i == 0:  # skip background
+            continue
+        metrics[class_names[i]] = ap
+        result_str += "{:<16}: {:.4f}\n".format(class_names[i], ap)
+    logger.info(result_str)
+
+    if iteration is not None:
+        result_path = os.path.join(output_dir, 'result_{:07d}.txt'.format(iteration))
+    else:
+        result_path = os.path.join(output_dir, 'result_{}.txt'.format(datetime.now().strftime('%Y-%m-%d_%H-%M-%S')))
+    with open(result_path, "w") as f:
+        f.write(result_str)
+
+    return dict(metrics=metrics)

ssd/data/datasets/evaluation/voc/eval_detection_voc.py

@@ -0,0 +1,346 @@
+from __future__ import division
+
+from collections import defaultdict
+import itertools
+import numpy as np
+import six
+
+
+def bbox_iou(bbox_a, bbox_b):
+    """Calculate the Intersection of Unions (IoUs) between bounding boxes.
+    IoU is calculated as a ratio of area of the intersection
+    and area of the union.
+    This function accepts both :obj:`numpy.ndarray` and :obj:`cupy.ndarray` as
+    inputs. Please note that both :obj:`bbox_a` and :obj:`bbox_b` need to be
+    same type.
+    The output is same type as the type of the inputs.
+    Args:
+        bbox_a (array): An array whose shape is :math:`(N, 4)`.
+            :math:`N` is the number of bounding boxes.
+            The dtype should be :obj:`numpy.float32`.
+        bbox_b (array): An array similar to :obj:`bbox_a`,
+            whose shape is :math:`(K, 4)`.
+            The dtype should be :obj:`numpy.float32`.
+    Returns:
+        array:
+        An array whose shape is :math:`(N, K)`. \
+        An element at index :math:`(n, k)` contains IoUs between \
+        :math:`n` th bounding box in :obj:`bbox_a` and :math:`k` th bounding \
+        box in :obj:`bbox_b`.
+    """
+    if bbox_a.shape[1] != 4 or bbox_b.shape[1] != 4:
+        raise IndexError
+
+    # top left
+    tl = np.maximum(bbox_a[:, None, :2], bbox_b[:, :2])
+    # bottom right
+    br = np.minimum(bbox_a[:, None, 2:], bbox_b[:, 2:])
+
+    area_i = np.prod(br - tl, axis=2) * (tl < br).all(axis=2)
+    area_a = np.prod(bbox_a[:, 2:] - bbox_a[:, :2], axis=1)
+    area_b = np.prod(bbox_b[:, 2:] - bbox_b[:, :2], axis=1)
+    return area_i / (area_a[:, None] + area_b - area_i)
+
+
+def eval_detection_voc(
+        pred_bboxes,
+        pred_labels,
+        pred_scores,
+        gt_bboxes,
+        gt_labels,
+        gt_difficults=None,
+        iou_thresh=0.5,
+        use_07_metric=False):
+    """Calculate average precisions based on evaluation code of PASCAL VOC.
+
+    This function evaluates predicted bounding boxes obtained from a dataset
+    which has :math:`N` images by using average precision for each class.
+    The code is based on the evaluation code used in PASCAL VOC Challenge.
+
+    Args:
+        pred_bboxes (iterable of numpy.ndarray): An iterable of :math:`N`
+            sets of bounding boxes.
+            Its index corresponds to an index for the base dataset.
+            Each element of :obj:`pred_bboxes` is a set of coordinates
+            of bounding boxes. This is an array whose shape is :math:`(R, 4)`,
+            where :math:`R` corresponds
+            to the number of bounding boxes, which may vary among boxes.
+            The second axis corresponds to
+            :math:`y_{min}, x_{min}, y_{max}, x_{max}` of a bounding box.
+        pred_labels (iterable of numpy.ndarray): An iterable of labels.
+            Similar to :obj:`pred_bboxes`, its index corresponds to an
+            index for the base dataset. Its length is :math:`N`.
+        pred_scores (iterable of numpy.ndarray): An iterable of confidence
+            scores for predicted bounding boxes. Similar to :obj:`pred_bboxes`,
+            its index corresponds to an index for the base dataset.
+            Its length is :math:`N`.
+        gt_bboxes (iterable of numpy.ndarray): An iterable of ground truth
+            bounding boxes
+            whose length is :math:`N`. An element of :obj:`gt_bboxes` is a
+            bounding box whose shape is :math:`(R, 4)`. Note that the number of
+            bounding boxes in each image does not need to be same as the number
+            of corresponding predicted boxes.
+        gt_labels (iterable of numpy.ndarray): An iterable of ground truth
+            labels which are organized similarly to :obj:`gt_bboxes`.
+        gt_difficults (iterable of numpy.ndarray): An iterable of boolean
+            arrays which is organized similarly to :obj:`gt_bboxes`.
+            This tells whether the
+            corresponding ground truth bounding box is difficult or not.
+            By default, this is :obj:`None`. In that case, this function
+            considers all bounding boxes to be not difficult.
+        iou_thresh (float): A prediction is correct if its Intersection over
+            Union with the ground truth is above this value.
+        use_07_metric (bool): Whether to use PASCAL VOC 2007 evaluation metric
+            for calculating average precision. The default value is
+            :obj:`False`.
+
+    Returns:
+        dict:
+
+        The keys, value-types and the description of the values are listed
+        below.
+
+        * **ap** (*numpy.ndarray*): An array of average precisions. \
+            The :math:`l`-th value corresponds to the average precision \
+            for class :math:`l`. If class :math:`l` does not exist in \
+            either :obj:`pred_labels` or :obj:`gt_labels`, the corresponding \
+            value is set to :obj:`numpy.nan`.
+        * **map** (*float*): The average of Average Precisions over classes.
+
+    """
+
+    prec, rec = calc_detection_voc_prec_rec(pred_bboxes,
+                                            pred_labels,
+                                            pred_scores,
+                                            gt_bboxes,
+                                            gt_labels,
+                                            gt_difficults,
+                                            iou_thresh=iou_thresh)
+
+    ap = calc_detection_voc_ap(prec, rec, use_07_metric=use_07_metric)
+
+    return {'ap': ap, 'map': np.nanmean(ap)}
+
+
+def calc_detection_voc_prec_rec(
+        pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels,
+        gt_difficults=None,
+        iou_thresh=0.5):
+    """Calculate precision and recall based on evaluation code of PASCAL VOC.
+
+    This function calculates precision and recall of
+    predicted bounding boxes obtained from a dataset which has :math:`N`
+    images.
+    The code is based on the evaluation code used in PASCAL VOC Challenge.
+
+    Args:
+        pred_bboxes (iterable of numpy.ndarray): An iterable of :math:`N`
+            sets of bounding boxes.
+            Its index corresponds to an index for the base dataset.
+            Each element of :obj:`pred_bboxes` is a set of coordinates
+            of bounding boxes. This is an array whose shape is :math:`(R, 4)`,
+            where :math:`R` corresponds
+            to the number of bounding boxes, which may vary among boxes.
+            The second axis corresponds to
+            :math:`y_{min}, x_{min}, y_{max}, x_{max}` of a bounding box.
+        pred_labels (iterable of numpy.ndarray): An iterable of labels.
+            Similar to :obj:`pred_bboxes`, its index corresponds to an
+            index for the base dataset. Its length is :math:`N`.
+        pred_scores (iterable of numpy.ndarray): An iterable of confidence
+            scores for predicted bounding boxes. Similar to :obj:`pred_bboxes`,
+            its index corresponds to an index for the base dataset.
+            Its length is :math:`N`.
+        gt_bboxes (iterable of numpy.ndarray): An iterable of ground truth
+            bounding boxes
+            whose length is :math:`N`. An element of :obj:`gt_bboxes` is a
+            bounding box whose shape is :math:`(R, 4)`. Note that the number of
+            bounding boxes in each image does not need to be same as the number
+            of corresponding predicted boxes.
+        gt_labels (iterable of numpy.ndarray): An iterable of ground truth
+            labels which are organized similarly to :obj:`gt_bboxes`.
+        gt_difficults (iterable of numpy.ndarray): An iterable of boolean
+            arrays which is organized similarly to :obj:`gt_bboxes`.
+            This tells whether the
+            corresponding ground truth bounding box is difficult or not.
+            By default, this is :obj:`None`. In that case, this function
+            considers all bounding boxes to be not difficult.
+        iou_thresh (float): A prediction is correct if its Intersection over
+            Union with the ground truth is above this value..
+
+    Returns:
+        tuple of two lists:
+        This function returns two lists: :obj:`prec` and :obj:`rec`.
+
+        * :obj:`prec`: A list of arrays. :obj:`prec[l]` is precision \
+            for class :math:`l`. If class :math:`l` does not exist in \
+            either :obj:`pred_labels` or :obj:`gt_labels`, :obj:`prec[l]` is \
+            set to :obj:`None`.
+        * :obj:`rec`: A list of arrays. :obj:`rec[l]` is recall \
+            for class :math:`l`. If class :math:`l` that is not marked as \
+            difficult does not exist in \
+            :obj:`gt_labels`, :obj:`rec[l]` is \
+            set to :obj:`None`.
+
+    """
+
+    pred_bboxes = iter(pred_bboxes)
+    pred_labels = iter(pred_labels)
+    pred_scores = iter(pred_scores)
+    gt_bboxes = iter(gt_bboxes)
+    gt_labels = iter(gt_labels)
+    if gt_difficults is None:
+        gt_difficults = itertools.repeat(None)
+    else:
+        gt_difficults = iter(gt_difficults)
+
+    n_pos = defaultdict(int)
+    score = defaultdict(list)
+    match = defaultdict(list)
+
+    for pred_bbox, pred_label, pred_score, gt_bbox, gt_label, gt_difficult in \
+            six.moves.zip(
+                pred_bboxes, pred_labels, pred_scores,
+                gt_bboxes, gt_labels, gt_difficults):
+
+        if gt_difficult is None:
+            gt_difficult = np.zeros(gt_bbox.shape[0], dtype=bool)
+
+        for l in np.unique(np.concatenate((pred_label, gt_label)).astype(int)):
+            pred_mask_l = pred_label == l
+            pred_bbox_l = pred_bbox[pred_mask_l]
+            pred_score_l = pred_score[pred_mask_l]
+            # sort by score
+            order = pred_score_l.argsort()[::-1]
+            pred_bbox_l = pred_bbox_l[order]
+            pred_score_l = pred_score_l[order]
+
+            gt_mask_l = gt_label == l
+            gt_bbox_l = gt_bbox[gt_mask_l]
+            gt_difficult_l = gt_difficult[gt_mask_l]
+
+            n_pos[l] += np.logical_not(gt_difficult_l).sum()
+            score[l].extend(pred_score_l)
+
+            if len(pred_bbox_l) == 0:
+                continue
+            if len(gt_bbox_l) == 0:
+                match[l].extend((0,) * pred_bbox_l.shape[0])
+                continue
+
+            # VOC evaluation follows integer typed bounding boxes.
+            pred_bbox_l = pred_bbox_l.copy()
+            pred_bbox_l[:, 2:] += 1
+            gt_bbox_l = gt_bbox_l.copy()
+            gt_bbox_l[:, 2:] += 1
+
+            iou = bbox_iou(pred_bbox_l, gt_bbox_l)
+            gt_index = iou.argmax(axis=1)
+            # set -1 if there is no matching ground truth
+            gt_index[iou.max(axis=1) < iou_thresh] = -1
+            del iou
+
+            selec = np.zeros(gt_bbox_l.shape[0], dtype=bool)
+            for gt_idx in gt_index:
+                if gt_idx >= 0:
+                    if gt_difficult_l[gt_idx]:
+                        match[l].append(-1)
+                    else:
+                        if not selec[gt_idx]:
+                            match[l].append(1)
+                        else:
+                            match[l].append(0)
+                    selec[gt_idx] = True
+                else:
+                    match[l].append(0)
+
+    for iter_ in (
+            pred_bboxes, pred_labels, pred_scores,
+            gt_bboxes, gt_labels, gt_difficults):
+        if next(iter_, None) is not None:
+            raise ValueError('Length of input iterables need to be same.')
+
+    n_fg_class = max(n_pos.keys()) + 1
+    prec = [None] * n_fg_class
+    rec = [None] * n_fg_class
+
+    for l in n_pos.keys():
+        score_l = np.array(score[l])
+        match_l = np.array(match[l], dtype=np.int8)
+
+        order = score_l.argsort()[::-1]
+        match_l = match_l[order]
+
+        tp = np.cumsum(match_l == 1)
+        fp = np.cumsum(match_l == 0)
+
+        # If an element of fp + tp is 0,
+        # the corresponding element of prec[l] is nan.
+        prec[l] = tp / (fp + tp)
+        # If n_pos[l] is 0, rec[l] is None.
+        if n_pos[l] > 0:
+            rec[l] = tp / n_pos[l]
+
+    return prec, rec
+
+
+def calc_detection_voc_ap(prec, rec, use_07_metric=False):
+    """Calculate average precisions based on evaluation code of PASCAL VOC.
+
+    This function calculates average precisions
+    from given precisions and recalls.
+    The code is based on the evaluation code used in PASCAL VOC Challenge.
+
+    Args:
+        prec (list of numpy.array): A list of arrays.
+            :obj:`prec[l]` indicates precision for class :math:`l`.
+            If :obj:`prec[l]` is :obj:`None`, this function returns
+            :obj:`numpy.nan` for class :math:`l`.
+        rec (list of numpy.array): A list of arrays.
+            :obj:`rec[l]` indicates recall for class :math:`l`.
+            If :obj:`rec[l]` is :obj:`None`, this function returns
+            :obj:`numpy.nan` for class :math:`l`.
+        use_07_metric (bool): Whether to use PASCAL VOC 2007 evaluation metric
+            for calculating average precision. The default value is
+            :obj:`False`.
+
+    Returns:
+        ~numpy.ndarray:
+        This function returns an array of average precisions.
+        The :math:`l`-th value corresponds to the average precision
+        for class :math:`l`. If :obj:`prec[l]` or :obj:`rec[l]` is
+        :obj:`None`, the corresponding value is set to :obj:`numpy.nan`.
+
+    """
+
+    n_fg_class = len(prec)
+    ap = np.empty(n_fg_class)
+    for l in six.moves.range(n_fg_class):
+        if prec[l] is None or rec[l] is None:
+            ap[l] = np.nan
+            continue
+
+        if use_07_metric:
+            # 11 point metric
+            ap[l] = 0
+            for t in np.arange(0., 1.1, 0.1):
+                if np.sum(rec[l] >= t) == 0:
+                    p = 0
+                else:
+                    p = np.max(np.nan_to_num(prec[l])[rec[l] >= t])
+                ap[l] += p / 11
+        else:
+            # correct AP calculation
+            # first append sentinel values at the end
+            mpre = np.concatenate(([0], np.nan_to_num(prec[l]), [0]))
+            mrec = np.concatenate(([0], rec[l], [1]))
+
+            mpre = np.maximum.accumulate(mpre[::-1])[::-1]
+
+            # to calculate area under PR curve, look for points
+            # where X axis (recall) changes value
+            i = np.where(mrec[1:] != mrec[:-1])[0]
+
+            # and sum (\Delta recall) * prec
+            ap[l] = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
+
+    return ap

ssd/data/datasets/voc.py

@@ -0,0 +1,123 @@
+import os
+import torch.utils.data
+import numpy as np
+import xml.etree.ElementTree as ET
+from PIL import Image
+
+from ssd.structures.container import Container
+
+root_dir = "/data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit"
+class VOCDataset(torch.utils.data.Dataset):
+
+    class_names = ('__background__',
+                   'aeroplane', 'bicycle', 'bird', 'boat',
+                   'bottle', 'bus', 'car', 'cat', 'chair',
+                   'cow', 'diningtable', 'dog', 'horse',
+                   'motorbike', 'person', 'pottedplant',
+                   'sheep', 'sofa', 'train', 'tvmonitor')
+
+    def __init__(self, data_dir, split, transform=None, target_transform=None, keep_difficult=False):
+        """Dataset for VOC data.
+        Args:
+            data_dir: the root of the VOC2007 or VOC2012 dataset, the directory contains the following sub-directories:
+                Annotations, ImageSets, JPEGImages, SegmentationClass, SegmentationObject.
+        """
+        self.data_dir = data_dir
+        self.split = split
+        self.transform = transform
+        self.target_transform = target_transform
+        image_sets_file = os.path.join(root_dir, self.data_dir, "ImageSets", "Main", "%s.txt" % self.split)
+        self.ids = VOCDataset._read_image_ids(image_sets_file)
+        self.keep_difficult = keep_difficult
+
+        self.class_dict = {class_name: i for i, class_name in enumerate(self.class_names)}
+
+    def __getitem__(self, index):
+        image_id = self.ids[index]
+        boxes, labels, is_difficult = self._get_annotation(image_id)
+        if not self.keep_difficult:
+            boxes = boxes[is_difficult == 0]
+            labels = labels[is_difficult == 0]
+        image = self._read_image(image_id)
+        if self.transform:
+            # print(image_id,"is dealing")
+            # print(boxes)
+            image, boxes, labels = self.transform(image, boxes, labels)
+        if self.target_transform:
+            boxes, labels = self.target_transform(boxes, labels)
+        targets = Container(
+            boxes=boxes,
+            labels=labels,
+        )
+        return image, targets, index
+    # def __getitem__(self, index):
+    #     image_id = self.ids[index]
+    #     #boxes, labels, is_difficult = self._get_annotation(image_id)
+    #   #  if not self.keep_difficult:
+    #       #  boxes = boxes[is_difficult == 0]
+    #        # labels = labels[is_difficult == 0]
+    #     image = self._read_image(image_id)
+    #     # if self.transform:
+    #     #     image, boxes, labels = self.transform(image, boxes, labels)
+    #     if self.transform:
+    #         image, boxes, labels = self.transform(image)
+    #     # if self.target_transform:
+    #     #     boxes, labels = self.target_transform(boxes, labels)
+    #     targets = Container(
+    #         boxes=boxes,
+    #         labels=labels,
+    #     )
+    #     return image, targets, index
+    def get_annotation(self, index):
+        image_id = self.ids[index]
+        return image_id, self._get_annotation(image_id)
+
+    def __len__(self):
+        return len(self.ids)
+
+    @staticmethod
+    def _read_image_ids(image_sets_file):
+        ids = []
+        with open(image_sets_file) as f:
+            for line in f:
+                ids.append(line.rstrip())
+        return ids
+
+    def _get_annotation(self, image_id):
+        annotation_file = os.path.join(root_dir, self.data_dir, "Annotations", "%s.xml" % image_id)
+        objects = ET.parse(annotation_file).findall("object")
+        boxes = []
+        labels = []
+        is_difficult = []
+        for obj in objects:
+            class_name = obj.find('name').text.lower().strip()
+            bbox = obj.find('bndbox')
+            # VOC dataset format follows Matlab, in which indexes start from 0
+            x1 = float(bbox.find('xmin').text) - 1
+            y1 = float(bbox.find('ymin').text) - 1
+            x2 = float(bbox.find('xmax').text) - 1
+            y2 = float(bbox.find('ymax').text) - 1
+            boxes.append([x1, y1, x2, y2])
+            labels.append(self.class_dict['bird'])
+           # labels.append(self.class_dict[class_name])
+            #labels.append('bird')
+            #is_difficult_str = obj.find('difficult').text
+            #is_difficult.append(int(is_difficult_str) if is_difficult_str else 0)
+            is_difficult.append(0)
+        return (np.array(boxes, dtype=np.float32),
+                np.array(labels, dtype=np.int64),
+                np.array(is_difficult, dtype=np.uint8))
+
+    def get_img_info(self, index):
+        img_id = self.ids[index]
+        annotation_file = os.path.join(root_dir, self.data_dir, "Annotations", "%s.xml" % img_id)
+        anno = ET.parse(annotation_file).getroot()
+        size = anno.find("size")
+        im_info = tuple(map(int, (size.find("height").text, size.find("width").text)))
+        return {"height": im_info[0], "width": im_info[1]}
+
+    def _read_image(self, image_id):
+        image_file = os.path.join(root_dir, self.data_dir, "JPEGImages", "%s.jpg" % image_id)
+        image = Image.open(image_file).convert("RGB")
+        image = np.array(image)
+        return image

ssd/data/samplers/__init__.py

@@ -0,0 +1,4 @@
+from .iteration_based_batch_sampler import IterationBasedBatchSampler
+from .distributed import DistributedSampler
+
+__all__ = ['IterationBasedBatchSampler', 'DistributedSampler']

ssd/data/samplers/distributed.py

@@ -0,0 +1,66 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# Code is copy-pasted exactly as in torch.utils.data.distributed.
+# FIXME remove this once c10d fixes the bug it has
+import math
+import torch
+import torch.distributed as dist
+from torch.utils.data.sampler import Sampler
+
+
+class DistributedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+
+    def __iter__(self):
+        if self.shuffle:
+            # deterministically shuffle based on epoch
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = torch.arange(len(self.dataset)).tolist()
+
+        # add extra samples to make it evenly divisible
+        indices += indices[: (self.total_size - len(indices))]
+        assert len(indices) == self.total_size
+
+        # subsample
+        offset = self.num_samples * self.rank
+        indices = indices[offset: offset + self.num_samples]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

ssd/data/samplers/iteration_based_batch_sampler.py

@@ -0,0 +1,30 @@
+from torch.utils.data.sampler import BatchSampler
+
+
+class IterationBasedBatchSampler(BatchSampler):
+    """
+    Wraps a BatchSampler, re-sampling from it until
+    a specified number of iterations have been sampled
+    """
+
+    def __init__(self, batch_sampler, num_iterations, start_iter=0):
+        self.batch_sampler = batch_sampler
+        self.num_iterations = num_iterations
+        self.start_iter = start_iter
+
+    def __iter__(self):
+        iteration = self.start_iter
+        while iteration <= self.num_iterations:
+            # if the underlying sampler has a set_epoch method, like
+            # DistributedSampler, used for making each process see
+            # a different split of the dataset, then set it
+            if hasattr(self.batch_sampler.sampler, "set_epoch"):
+                self.batch_sampler.sampler.set_epoch(iteration)
+            for batch in self.batch_sampler:
+                iteration += 1
+                if iteration > self.num_iterations:
+                    break
+                yield batch
+
+    def __len__(self):
+        return self.num_iterations

ssd/data/transforms/__init__.py

@@ -0,0 +1,34 @@
+from ssd.modeling.anchors.prior_box import PriorBox
+from .target_transform import SSDTargetTransform
+from .transforms import *
+
+
+def build_transforms(cfg, is_train=True):
+    if is_train:
+        transform = [
+            ConvertFromInts(),
+            PhotometricDistort(),
+            Expand(cfg.INPUT.PIXEL_MEAN),
+            RandomSampleCrop(),
+            RandomMirror(),
+            ToPercentCoords(),
+            Resize(cfg.INPUT.IMAGE_SIZE),
+            SubtractMeans(cfg.INPUT.PIXEL_MEAN),
+            ToTensor(),
+        ]
+    else:
+        transform = [
+            Resize(cfg.INPUT.IMAGE_SIZE),
+            SubtractMeans(cfg.INPUT.PIXEL_MEAN),
+            ToTensor()
+        ]
+    transform = Compose(transform)
+    return transform
+
+
+def build_target_transform(cfg):
+    transform = SSDTargetTransform(PriorBox(cfg)(),
+                                   cfg.MODEL.CENTER_VARIANCE,
+                                   cfg.MODEL.SIZE_VARIANCE,
+                                   cfg.MODEL.THRESHOLD)
+    return transform

ssd/data/transforms/target_transform.py

@@ -0,0 +1,25 @@
+import numpy as np
+import torch
+
+from ssd.utils import box_utils
+
+
+class SSDTargetTransform:
+    def __init__(self, center_form_priors, center_variance, size_variance, iou_threshold):
+        self.center_form_priors = center_form_priors
+        self.corner_form_priors = box_utils.center_form_to_corner_form(center_form_priors)
+        self.center_variance = center_variance
+        self.size_variance = size_variance
+        self.iou_threshold = iou_threshold
+
+    def __call__(self, gt_boxes, gt_labels):
+        if type(gt_boxes) is np.ndarray:
+            gt_boxes = torch.from_numpy(gt_boxes)
+        if type(gt_labels) is np.ndarray:
+            gt_labels = torch.from_numpy(gt_labels)
+        boxes, labels = box_utils.assign_priors(gt_boxes, gt_labels,
+                                                self.corner_form_priors, self.iou_threshold)
+        boxes = box_utils.corner_form_to_center_form(boxes)
+        locations = box_utils.convert_boxes_to_locations(boxes, self.center_form_priors, self.center_variance, self.size_variance)
+       
+        return locations, labels

ssd/data/transforms/transforms.py

@@ -0,0 +1,410 @@
+# from https://github.com/amdegroot/ssd.pytorch
+
+
+import torch
+from torchvision import transforms
+import cv2
+import numpy as np
+import types
+from numpy import random
+
+
+def intersect(box_a, box_b):
+    max_xy = np.minimum(box_a[:, 2:], box_b[2:])
+    min_xy = np.maximum(box_a[:, :2], box_b[:2])
+    inter = np.clip((max_xy - min_xy), a_min=0, a_max=np.inf)
+    return inter[:, 0] * inter[:, 1]
+
+
+def jaccard_numpy(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: Multiple bounding boxes, Shape: [num_boxes,4]
+        box_b: Single bounding box, Shape: [4]
+    Return:
+        jaccard overlap: Shape: [box_a.shape[0], box_a.shape[1]]
+    """
+    inter = intersect(box_a, box_b)
+    area_a = ((box_a[:, 2] - box_a[:, 0]) *
+              (box_a[:, 3] - box_a[:, 1]))  # [A,B]
+    area_b = ((box_b[2] - box_b[0]) *
+              (box_b[3] - box_b[1]))  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+class Compose(object):
+    """Composes several augmentations together.
+    Args:
+        transforms (List[Transform]): list of transforms to compose.
+    Example:
+        >>> augmentations.Compose([
+        >>>     transforms.CenterCrop(10),
+        >>>     transforms.ToTensor(),
+        >>> ])
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img, boxes=None, labels=None):
+        for t in self.transforms:
+            img, boxes, labels = t(img, boxes, labels)
+        return img, boxes, labels
+
+
+class Lambda(object):
+    """Applies a lambda as a transform."""
+
+    def __init__(self, lambd):
+        assert isinstance(lambd, types.LambdaType)
+        self.lambd = lambd
+
+    def __call__(self, img, boxes=None, labels=None):
+        return self.lambd(img, boxes, labels)
+
+
+class ConvertFromInts(object):
+    def __call__(self, image, boxes=None, labels=None):
+        return image.astype(np.float32), boxes, labels
+
+
+class SubtractMeans(object):
+    def __init__(self, mean):
+        self.mean = np.array(mean, dtype=np.float32)
+
+    def __call__(self, image, boxes=None, labels=None):
+        image = image.astype(np.float32)
+        image -= self.mean
+        return image.astype(np.float32), boxes, labels
+
+
+class ToAbsoluteCoords(object):
+    def __call__(self, image, boxes=None, labels=None):
+        height, width, channels = image.shape
+        boxes[:, 0] *= width
+        boxes[:, 2] *= width
+        boxes[:, 1] *= height
+        boxes[:, 3] *= height
+
+        return image, boxes, labels
+
+
+class ToPercentCoords(object):
+    def __call__(self, image, boxes=None, labels=None):
+        height, width, channels = image.shape
+        boxes[:, 0] /= width
+        boxes[:, 2] /= width
+        boxes[:, 1] /= height
+        boxes[:, 3] /= height
+
+        return image, boxes, labels
+
+
+class Resize(object):
+    def __init__(self, size=300):
+        self.size = size
+
+    def __call__(self, image, boxes=None, labels=None):
+        image = cv2.resize(image, (self.size,
+                                   self.size))
+        return image, boxes, labels
+
+
+class RandomSaturation(object):
+    def __init__(self, lower=0.5, upper=1.5):
+        self.lower = lower
+        self.upper = upper
+        assert self.upper >= self.lower, "contrast upper must be >= lower."
+        assert self.lower >= 0, "contrast lower must be non-negative."
+
+    def __call__(self, image, boxes=None, labels=None):
+        if random.randint(2):
+            image[:, :, 1] *= random.uniform(self.lower, self.upper)
+
+        return image, boxes, labels
+
+
+class RandomHue(object):
+    def __init__(self, delta=18.0):
+        assert delta >= 0.0 and delta <= 360.0
+        self.delta = delta
+
+    def __call__(self, image, boxes=None, labels=None):
+        if random.randint(2):
+            image[:, :, 0] += random.uniform(-self.delta, self.delta)
+            image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0
+            image[:, :, 0][image[:, :, 0] < 0.0] += 360.0
+        return image, boxes, labels
+
+
+class RandomLightingNoise(object):
+    def __init__(self):
+        self.perms = ((0, 1, 2), (0, 2, 1),
+                      (1, 0, 2), (1, 2, 0),
+                      (2, 0, 1), (2, 1, 0))
+
+    def __call__(self, image, boxes=None, labels=None):
+        if random.randint(2):
+            swap = self.perms[random.randint(len(self.perms))]
+            shuffle = SwapChannels(swap)  # shuffle channels
+            image = shuffle(image)
+        return image, boxes, labels
+
+
+class ConvertColor(object):
+    def __init__(self, current, transform):
+        self.transform = transform
+        self.current = current
+
+    def __call__(self, image, boxes=None, labels=None):
+        if self.current == 'BGR' and self.transform == 'HSV':
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+        elif self.current == 'RGB' and self.transform == 'HSV':
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
+        elif self.current == 'BGR' and self.transform == 'RGB':
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        elif self.current == 'HSV' and self.transform == 'BGR':
+            image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+        elif self.current == 'HSV' and self.transform == "RGB":
+            image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
+        else:
+            raise NotImplementedError
+        return image, boxes, labels
+
+
+class RandomContrast(object):
+    def __init__(self, lower=0.5, upper=1.5):
+        self.lower = lower
+        self.upper = upper
+        assert self.upper >= self.lower, "contrast upper must be >= lower."
+        assert self.lower >= 0, "contrast lower must be non-negative."
+
+    # expects float image
+    def __call__(self, image, boxes=None, labels=None):
+        if random.randint(2):
+            alpha = random.uniform(self.lower, self.upper)
+            image *= alpha
+        return image, boxes, labels
+
+
+class RandomBrightness(object):
+    def __init__(self, delta=32):
+        assert delta >= 0.0
+        assert delta <= 255.0
+        self.delta = delta
+
+    def __call__(self, image, boxes=None, labels=None):
+        if random.randint(2):
+            delta = random.uniform(-self.delta, self.delta)
+            image += delta
+        return image, boxes, labels
+
+
+class ToCV2Image(object):
+    def __call__(self, tensor, boxes=None, labels=None):
+        return tensor.cpu().numpy().astype(np.float32).transpose((1, 2, 0)), boxes, labels
+
+
+class ToTensor(object):
+    def __call__(self, cvimage, boxes=None, labels=None):
+        return torch.from_numpy(cvimage.astype(np.float32)).permute(2, 0, 1), boxes, labels
+
+
+class RandomSampleCrop(object):
+    """Crop
+    Arguments:
+        img (Image): the image being input during training
+        boxes (Tensor): the original bounding boxes in pt form
+        labels (Tensor): the class labels for each bbox
+        mode (float tuple): the min and max jaccard overlaps
+    Return:
+        (img, boxes, classes)
+            img (Image): the cropped image
+            boxes (Tensor): the adjusted bounding boxes in pt form
+            labels (Tensor): the class labels for each bbox
+    """
+
+    def __init__(self):
+        self.sample_options = (
+            # using entire original input image
+            None,
+            # sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9
+            (0.1, None),
+            (0.3, None),
+            (0.7, None),
+            (0.9, None),
+            # randomly sample a patch
+            (None, None),
+        )
+
+    def __call__(self, image, boxes=None, labels=None):
+        # guard against no boxes
+        if boxes is not None and boxes.shape[0] == 0:
+            return image, boxes, labels
+        height, width, _ = image.shape
+        while True:
+            # randomly choose a mode
+            mode = random.choice(self.sample_options)
+            if mode is None:
+                return image, boxes, labels
+
+            min_iou, max_iou = mode
+            if min_iou is None:
+                min_iou = float('-inf')
+            if max_iou is None:
+                max_iou = float('inf')
+
+            # max trails (50)
+            for _ in range(50):
+                current_image = image
+
+                w = random.uniform(0.3 * width, width)
+                h = random.uniform(0.3 * height, height)
+
+                # aspect ratio constraint b/t .5 & 2
+                if h / w < 0.5 or h / w > 2:
+                    continue
+
+                left = random.uniform(width - w)
+                top = random.uniform(height - h)
+
+                # convert to integer rect x1,y1,x2,y2
+                rect = np.array([int(left), int(top), int(left + w), int(top + h)])
+
+                # calculate IoU (jaccard overlap) b/t the cropped and gt boxes
+                overlap = jaccard_numpy(boxes, rect)
+
+                # is min and max overlap constraint satisfied? if not try again
+                if overlap.max() < min_iou or overlap.min() > max_iou:
+                    continue
+
+                # cut the crop from the image
+                current_image = current_image[rect[1]:rect[3], rect[0]:rect[2],
+                                :]
+
+                # keep overlap with gt box IF center in sampled patch
+                centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0
+
+                # mask in all gt boxes that above and to the left of centers
+                m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1])
+
+                # mask in all gt boxes that under and to the right of centers
+                m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1])
+
+                # mask in that both m1 and m2 are true
+                mask = m1 * m2
+
+                # have any valid boxes? try again if not
+                if not mask.any():
+                    continue
+
+                # take only matching gt boxes
+                current_boxes = boxes[mask, :].copy()
+
+                # take only matching gt labels
+                current_labels = labels[mask]
+
+                # should we use the box left and top corner or the crop's
+                current_boxes[:, :2] = np.maximum(current_boxes[:, :2],
+                                                  rect[:2])
+                # adjust to crop (by substracting crop's left,top)
+                current_boxes[:, :2] -= rect[:2]
+
+                current_boxes[:, 2:] = np.minimum(current_boxes[:, 2:],
+                                                  rect[2:])
+                # adjust to crop (by substracting crop's left,top)
+                current_boxes[:, 2:] -= rect[:2]
+
+                return current_image, current_boxes, current_labels
+
+
+class Expand(object):
+    def __init__(self, mean):
+        self.mean = mean
+
+    def __call__(self, image, boxes, labels):
+        if random.randint(2):
+            return image, boxes, labels
+
+        height, width, depth = image.shape
+        ratio = random.uniform(1, 4)
+        left = random.uniform(0, width * ratio - width)
+        top = random.uniform(0, height * ratio - height)
+
+        expand_image = np.zeros(
+            (int(height * ratio), int(width * ratio), depth),
+            dtype=image.dtype)
+        expand_image[:, :, :] = self.mean
+        expand_image[int(top):int(top + height),
+        int(left):int(left + width)] = image
+        image = expand_image
+
+        boxes = boxes.copy()
+        boxes[:, :2] += (int(left), int(top))
+        boxes[:, 2:] += (int(left), int(top))
+
+        return image, boxes, labels
+
+
+class RandomMirror(object):
+    def __call__(self, image, boxes, classes):
+        _, width, _ = image.shape
+        if random.randint(2):
+            image = image[:, ::-1]
+            boxes = boxes.copy()
+            boxes[:, 0::2] = width - boxes[:, 2::-2]
+        return image, boxes, classes
+
+
+class SwapChannels(object):
+    """Transforms a tensorized image by swapping the channels in the order
+     specified in the swap tuple.
+    Args:
+        swaps (int triple): final order of channels
+            eg: (2, 1, 0)
+    """
+
+    def __init__(self, swaps):
+        self.swaps = swaps
+
+    def __call__(self, image):
+        """
+        Args:
+            image (Tensor): image tensor to be transformed
+        Return:
+            a tensor with channels swapped according to swap
+        """
+        # if torch.is_tensor(image):
+        #     image = image.data.cpu().numpy()
+        # else:
+        #     image = np.array(image)
+        image = image[:, :, self.swaps]
+        return image
+
+
+class PhotometricDistort(object):
+    def __init__(self):
+        self.pd = [
+            RandomContrast(),  # RGB
+            ConvertColor(current="RGB", transform='HSV'),  # HSV
+            RandomSaturation(),  # HSV
+            RandomHue(),  # HSV
+            ConvertColor(current='HSV', transform='RGB'),  # RGB
+            RandomContrast()  # RGB
+        ]
+        self.rand_brightness = RandomBrightness()
+        self.rand_light_noise = RandomLightingNoise()
+
+    def __call__(self, image, boxes, labels):
+        im = image.copy()
+        im, boxes, labels = self.rand_brightness(im, boxes, labels)
+        if random.randint(2):
+            distort = Compose(self.pd[:-1])
+        else:
+            distort = Compose(self.pd[1:])
+        im, boxes, labels = distort(im, boxes, labels)
+        return self.rand_light_noise(im, boxes, labels)

ssd/engine/inference.py

@@ -0,0 +1,422 @@
+import logging
+import os
+from vizer.draw import draw_boxes
+import torch
+import torch.utils.data
+from tqdm import tqdm
+from ssd.config import cfg
+from ssd.data.build import make_data_loader
+from ssd.data.datasets.evaluation import evaluate
+import datetime
+from ssd.utils import dist_util, mkdir
+from ssd.utils.dist_util import synchronize, is_main_process
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+import logging
+import time
+import os
+import cv2
+import torch
+from tqdm import tqdm
+from ssd.utils.dist_util import is_main_process, get_world_size
+from ssd.utils.dist_util import synchronize
+from ssd.utils.timer import Timer, get_time_str
+
+import numpy as np
+from PIL import Image
+
+def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
+    all_predictions = dist_util.all_gather(predictions_per_gpu)
+    if not dist_util.is_main_process():
+        return
+    # merge the list of dicts
+    predictions = {}
+    for p in all_predictions:
+        predictions.update(p)
+    # convert a dict where the key is the index in a list
+    image_ids = list(sorted(predictions.keys()))
+    if len(image_ids) != image_ids[-1] + 1:
+        logger = logging.getLogger("SSD.inference")
+        logger.warning(
+            "Number of images that were gathered from multiple processes is not "
+            "a contiguous set. Some images might be missing from the evaluation"
+        )
+
+    # convert to a list
+    predictions = [predictions[i] for i in image_ids]
+    return predictions
+
+
+def compute_on_dataset_byHatcherW(model, data_loader, device,timer=None):
+    results_dict = {}
+    cpu_device = torch.device("cpu")
+    now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+    filename = os.path.join(
+        "/home/baai/voc/HatcherW/SSD/", now_time)
+    isExists = os.path.exists(filename)
+    if not isExists:
+        os.makedirs(filename)
+        print(filename + ' 创建成功')
+    else:
+        print(filename + ' 目录已存在')
+    with open(
+            "/home/baai/voc/HatcherW/ImageSets/Main/test.txt",
+            "r") as f:
+        for batch in tqdm(data_loader):
+            # 每次进循环先读取一次文件路径
+
+
+            JPEG_NAME = f.readline().strip('\n')
+            JPEG_NAME = JPEG_NAME + '.jpg'
+            print(JPEG_NAME)
+            filename = os.path.join(
+                "/home/baai/voc/HatcherW/JPEGImages/",
+                JPEG_NAME)
+            print(filename)
+####################################
+            # src = cv2.imread(filename,flags=1)
+            # cv2.namedWindow("yuanlai", cv2.WINDOW_NORMAL)
+            # cv2.imshow("yuanlai", src)
+            # contrast_brightness_image(src, 0.5, 5)  # 第一个1.2为对比度  第二个为亮度数值越大越亮
+
+
+
+########################################
+
+
+            image_forinf = Image.open(filename).convert("RGB")  #read image
+            image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]] #conver image into array
+           # cv2.imshow('测试', image_forinf)
+            images, target, image_ids = batch  # target in inference does not be used
+
+
+            with torch.no_grad():
+                if timer:
+                    timer.tic()
+                    output = model(images.to(device))
+                if timer:
+                    if not cfg.MODEL.DEVICE == 'cpu':
+                        torch.cuda.synchronize()
+                    timer.toc()
+
+                output = [o.to(cpu_device) for o in output]
+                ##########inference 测试,用于测试如何提取数据
+                print("output is {}",format(output))
+                output_temp = output[0]
+                # print("output 类型 is {}".format(type(output[0]['scores'][0])))
+                # print("output fenshu is {}".format((output[0]['scores'][0])))
+                # print("output biankuang is {}".format((output_temp['boxes'][0,:])))
+                # print("output biankuangnum is {}".format((output_temp['boxes'].size())))
+                #print("output label is {}".format((output_temp['labels'][0])))
+                # print("hhhh")
+                bbox_nums = output_temp['boxes'].size()[0] #取出box的个数
+                output_scores = output[0]['scores'].numpy()
+                output_labels = output[0]['labels'].numpy()
+                output_boxes = output[0]['boxes'].numpy()
+
+                print("output 类型 is {}".format(type(output_scores)))
+                result = image_forinf.copy()
+                result1 = image_forinf.copy()
+                heigh = image_forinf.shape[0]
+                kuan  = image_forinf.shape[1]
+                print(heigh,kuan)
+                ###visual by hcwang
+                for bbox_num in range(bbox_nums):  # 画框  需要画bbox_nums个
+                    left_top = (
+                    (int)(output_boxes[bbox_num][0] * kuan / 300.0), (int)(output_boxes[bbox_num][1] * heigh / 300.0))
+                    right_bottom = (
+                    (int)(output_boxes[bbox_num][2] * kuan / 300.0), (int)(output_boxes[bbox_num][3] * heigh / 300.0))
+                    cv2.rectangle(result, left_top, right_bottom, (0, 0, 255), 3)
+                #cv2.imshow('测试',image_forinf)
+                # 存储
+                filename = os.path.join("/home/baai/voc/HatcherW/SSD/", now_time ,JPEG_NAME)
+              #  cv2.imwrite(filename, result)
+                for bbox_num in range(bbox_nums):
+                    output_boxes[bbox_num][0] *=kuan / 300.0
+                    output_boxes[bbox_num][1] *=heigh / 300.0
+                    output_boxes[bbox_num][2] *=kuan / 300.0
+                    output_boxes[bbox_num][3] *=heigh / 300.0
+                class_names = ('__background__',
+                   'aeroplane', 'bicycle', 'bird', 'boat',
+                   'bottle', 'bus', 'car', 'cat', 'chair',
+                   'cow', 'diningtable', 'dog', 'horse',
+                   'motorbike', 'person', 'pottedplant',
+                   'sheep', 'sofa', 'train', 'tvmonitor')
+                drawn_image = draw_boxes(result1, output_boxes, output_labels, output_scores, class_names).astype(np.uint8)
+                Image.fromarray(drawn_image).save(filename)
+                # cv2.imshow('result',drawn_image)   #显示结果
+                # cv2.waitKey(0)
+                # cv2.destroyAllWindows()
+
+                print(output)
+            results_dict.update(
+                {img_id: result for img_id, result in zip(image_ids, output)}
+            )
+    f.close()
+    return results_dict
+def compute_on_dataset(model, data_loader, device):
+    results_dict = {}
+    for batch in tqdm(data_loader):
+        images, targets, image_ids = batch
+        cpu_device = torch.device("cpu")
+        with torch.no_grad():
+            outputs = model(images.to(device))
+
+            outputs = [o.to(cpu_device) for o in outputs]
+        results_dict.update(
+            {img_id: result for img_id, result in zip(image_ids, outputs)}
+        )
+    return results_dict
+
+def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
+    dataset = data_loader.dataset
+    logger = logging.getLogger("SSD.inference")
+    logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
+    total_timer = Timer()
+    inference_timer = Timer()
+    total_timer.tic()
+
+    predictions_path = os.path.join(output_folder, 'predictions.pth')
+    if use_cached and os.path.exists(predictions_path):
+        predictions = torch.load(predictions_path, map_location='cpu')
+    else:
+        predictions = compute_on_dataset_byHatcherW(model, data_loader, device,inference_timer)
+        synchronize()
+        predictions = _accumulate_predictions_from_multiple_gpus(predictions)
+    if not is_main_process():
+        return
+    if output_folder:
+        torch.save(predictions, predictions_path)
+    #return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
+# def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
+#     dataset = data_loader.dataset
+#     logger = logging.getLogger("SSD.inference")
+#     logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
+#     total_timer = Timer()
+#     inference_timer = Timer()
+#     total_timer.tic()
+#     predictions_path = os.path.join(output_folder, 'predictions.pth')
+#     if use_cached and os.path.exists(predictions_path):
+#         predictions = torch.load(predictions_path, map_location='cpu')
+#     else:
+#         predictions = compute_on_dataset(model, data_loader, device,inference_timer)
+#         synchronize()
+#         predictions = _accumulate_predictions_from_multiple_gpus(predictions)
+#     if not is_main_process():
+#         return
+#     if output_folder:
+#         torch.save(predictions, predictions_path)
+#     return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
+
+def inference(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
+    dataset = data_loader.dataset
+    logger = logging.getLogger("SSD.inference")
+    logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
+    predictions_path = os.path.join(output_folder, 'predictions.pth')
+    if use_cached and os.path.exists(predictions_path):
+        predictions = torch.load(predictions_path, map_location='cpu')
+    else:
+        predictions = compute_on_dataset(model, data_loader, device)
+        synchronize()
+        predictions = _accumulate_predictions_from_multiple_gpus(predictions)
+    if not is_main_process():
+        return
+    if output_folder:
+        torch.save(predictions, predictions_path)
+    return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
+@torch.no_grad()
+def do_evaluation_byHatcherW(cfg, model, distributed, **kwargs):
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model = model.module
+    model.eval()
+    device = torch.device(cfg.MODEL.DEVICE)
+    data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
+    eval_results = []
+    for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
+        output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
+        if not os.path.exists(output_folder):
+            mkdir(output_folder)
+        eval_result = inference_byHatcherW(model, data_loader, dataset_name, device, output_folder, **kwargs)
+        eval_results.append(eval_result)
+
+
+def do_evaluation(cfg, model, distributed, **kwargs):
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model = model.module
+    model.eval()
+    device = torch.device(cfg.MODEL.DEVICE)
+    data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
+    eval_results = []
+    for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
+        output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
+        if not os.path.exists(output_folder):
+            mkdir(output_folder)
+        eval_result = inference(model, data_loader, dataset_name, device, output_folder, **kwargs)
+        eval_results.append(eval_result)
+    return eval_results
+# def inference_hcw(
+#         model,
+#         data_loader,
+#         dataset_name,
+#         iou_types=("bbox",),
+#         box_only=False,
+#         device="cuda",
+#         expected_results=(),
+#         expected_results_sigma_tol=4,
+#         output_folder=None,
+# ):
+#     # convert to a torch.device for efficiency
+#     device = torch.device(device)   #use GPU
+#     num_devices = get_world_size() #get the num of GPU
+#     logger = logging.getLogger("maskrcnn_benchmark.inference")
+#     dataset = data_loader.dataset #get dataset
+#     logger.info("Start evaluation on {} dataset({} images).".format(dataset_name, len(dataset)))
+#     total_timer = Timer()
+#     inference_timer = Timer()
+#     total_timer.tic()
+#     predictions = compute_on_dataset(model, data_loader, device, inference_timer)
+#     # wait for all processes to complete before measuring the time
+#     synchronize()
+#     total_time = total_timer.toc()
+#     total_time_str = get_time_str(total_time)
+#     logger.info(
+#         "Total run time: {} ({} s / img per device, on {} devices)".format(
+#             total_time_str, total_time * num_devices / len(dataset), num_devices
+#         )
+#     )
+#     total_infer_time = get_time_str(inference_timer.total_time)
+#     logger.info(
+#         "Model inference time: {} ({} s / img per device, on {} devices)".format(
+#             total_infer_time,
+#             inference_timer.total_time * num_devices / len(dataset),
+#             num_devices,
+#         )
+#     )
+#
+#     predictions = _accumulate_predictions_from_multiple_gpus(predictions)
+#     #visual_whc(predictions)
+#     if not is_main_process():
+#         return
+#
+#     if output_folder:
+#         torch.save(predictions, os.path.join(output_folder, "predictions.txt"))
+#
+#     extra_args = dict(
+#         box_only=box_only,
+#         iou_types=iou_types,
+#         expected_results=expected_results,
+#         expected_results_sigma_tol=expected_results_sigma_tol,
+#     )
+# def compute_on_dataset_hcw(model, data_loader, device, timer=None):
+#     model.eval()  #不启用BN和DP
+#     results_dict = {}   #the result
+#     cpu_device = torch.device("cpu")
+#     with open(
+#              "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
+#              "r") as f:
+#         for _, batch in enumerate(tqdm(data_loader)):
+#             #每次进循环先读取一次文件路径
+#             JPEG_NAME = f.readline().strip('\n')
+#             JPEG_NAME = JPEG_NAME + '.jpg'
+#             print(JPEG_NAME)
+#             filename = os.path.join(
+#                 "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/",
+#                 JPEG_NAME)
+#             print(filename)
+#             image_forinf = Image.open(filename).convert("RGB")
+#             image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]]
+#            # print("batch is "+batch)
+#             images, image_ids = batch
+#           #  print("image_id is "+image_ids)
+#             with torch.no_grad():  #确定是否反向传播求dao
+#                 if timer:
+#                     timer.tic()
+#                     output = model(images.to(device))
+#                 if timer:
+#                     if not cfg.MODEL.DEVICE == 'cpu':
+#                         torch.cuda.synchronize()
+#                     timer.toc()
+#                 output = [o.to(cpu_device) for o in output]
+#                 output_forinf = output[0]
+#                 print("output caculate success")
+#                 #将结果rezize
+#                 height, width = image_forinf.shape[:-1]
+#                 output_forinf = output_forinf.resize((width, height))
+#                 print("outout resize success")
+#
+#                 ################
+#                 top_predictions = select_top_predictions(output_forinf)
+#                 print("nms success")
+#                 result = image_forinf.copy()
+#                 result = overlay_boxes(result, top_predictions)
+#                 result = overlay_class_names(result, top_predictions)
+#                 print("print reslut handle success")
+#                 ################
+#                 filename = os.path.join(
+#                 "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result1/", JPEG_NAME)
+#                 cv2.imwrite(filename, result)
+#             results_dict.update(
+#                 {img_id: result for img_id, result in zip(image_ids, output)}
+#             )
+#             ##为了测试num_bboxex
+#             # print("output1:")
+#             # temp = results_dict[0].bbox.size()
+#             # print(temp[0])
+#             # print(temp[1])
+#             # print(results_dict[0].bbox.size())
+#             #
+#             #
+#             # print(type(results_dict[0].bbox))
+#
+#             # with open(
+#             #         "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
+#             #         "r") as f:
+#             #     JPEG_NAME = f.readline().strip('\n')
+#             #     #  JPEG_NAME.strip('\n')
+#             #     JPEG_NAME = JPEG_NAME + '.jpg'
+#             #     print(JPEG_NAME)
+#             # f.close()
+#             # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/", JPEG_NAME)
+#             # print(filename)
+#             # img = cv2.imread(filename)
+#             # #cv2.imshow('测试', img)
+#             # #cv2.waitKey(0)
+#             # #cv2.destroyAllWindows()  # dv2.destroyWindow(wname)
+#             # results_dict_np = results_dict[0].bbox.numpy()
+#             # left_top = ((int)(results_dict_np[0][0]*1920.0/1333.0), (int)(results_dict_np[0][1]*1920.0/1333.0))
+#             # right_bottom = ((int)(results_dict_np[0][2]*1920.0/1333.0), (int)(results_dict_np[0][3]*1920.0/1333.0))
+#             # cv2.rectangle(img, left_top, right_bottom, (0, 0, 255), 3)
+#             # print(results_dict_np[2][1])
+#             # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result/", JPEG_NAME)
+#             # cv2.imwrite(filename, img)
+#             print('hhhh')
+#             print('hhhh')
+#         #visual_whc(results_dict)
+#         f.close()
+#         return results_dict
+def do_evaluation_hcw(cfg, model, distributed, **kwargs):
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model = model.module
+    model.eval()
+    device = torch.device(cfg.MODEL.DEVICE)
+    data_loaders_vals = make_data_loader(cfg, is_train=False, distributed=distributed)
+    eval_results = []
+    output_folders = [None] * len(cfg.DATASETS.TEST)
+    if cfg.OUTPUT_DIR:
+        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
+            mkdir(output_folder)
+            output_folders[idx] = output_folder
+    for output_folder, data_loaders_val in zip(output_folders,cfg.DATASETS.TEST, data_loaders_vals):
+        eval_result = inference(model, data_loaders_val, dataset_name, device, output_folder, **kwargs)
+        eval_results.append(eval_result)
+    return eval_results
+
+
+def contrast_brightness_image(src1, a, g):
+    h, w, ch = src1.shape  # 获取shape的数值，height和width、通道
+
+    # 新建全零图片数组src2,将height和width，类型设置为原图片的通道类型(色素全为零，输出为全黑图片)
+    src2 = np.zeros([h, w, ch], src1.dtype)
+    dst = cv2.addWeighted(src1, a, src2, 1 - a, g)  # addWeighted函数说明如下
+    cv2.imshow("con-bri-demo", dst)
+

ssd/engine/trainer.py

@@ -0,0 +1,140 @@
+import collections
+import datetime
+import logging
+import os
+import time
+import torch
+import torch.distributed as dist
+
+from ssd.engine.inference import do_evaluation
+from ssd.utils import dist_util
+from ssd.utils.metric_logger import MetricLogger
+
+
+def write_metric(eval_result, prefix, summary_writer, global_step):
+    for key in eval_result:
+        value = eval_result[key]
+        tag = '{}/{}'.format(prefix, key)
+        if isinstance(value, collections.Mapping):
+            write_metric(value, tag, summary_writer, global_step)
+        else:
+            summary_writer.add_scalar(tag, value, global_step=global_step)
+
+
+def reduce_loss_dict(loss_dict):
+    """
+    Reduce the loss dictionary from all processes so that process with rank
+    0 has the averaged results. Returns a dict with the same fields as
+    loss_dict, after reduction.
+    """
+    world_size = dist_util.get_world_size()
+    if world_size < 2:
+        return loss_dict
+    with torch.no_grad():
+        loss_names = []
+        all_losses = []
+        for k in sorted(loss_dict.keys()):
+            loss_names.append(k)
+            all_losses.append(loss_dict[k])
+        all_losses = torch.stack(all_losses, dim=0)
+        dist.reduce(all_losses, dst=0)
+        if dist.get_rank() == 0:
+            # only main process gets accumulated, so only divide by
+            # world_size in this case
+            all_losses /= world_size
+        reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
+    return reduced_losses
+
+
+def do_train(cfg, model,
+             data_loader,
+             optimizer,
+             scheduler,
+             checkpointer,
+             device,
+             arguments,
+             args):
+    logger = logging.getLogger("SSD.trainer")
+    logger.info("Start training ...")
+    meters = MetricLogger()
+
+    model.train()
+    save_to_disk = dist_util.get_rank() == 0
+    if args.use_tensorboard and save_to_disk:
+        import tensorboardX
+
+        summary_writer = tensorboardX.SummaryWriter(log_dir=os.path.join(cfg.OUTPUT_DIR, 'tf_logs'))
+    else:
+        summary_writer = None
+
+    max_iter = len(data_loader)
+    start_iter = arguments["iteration"]
+    start_training_time = time.time()
+    end = time.time()
+    for iteration, (images, targets, _) in enumerate(data_loader, start_iter):
+        iteration = iteration + 1
+        arguments["iteration"] = iteration
+        scheduler.step()
+
+        images = images.to(device)
+        targets = targets.to(device)
+        loss_dict = model(images, targets=targets)
+        loss = sum(loss for loss in loss_dict.values())
+
+        # reduce losses over all GPUs for logging purposes
+        loss_dict_reduced = reduce_loss_dict(loss_dict)
+        losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+        meters.update(total_loss=losses_reduced, **loss_dict_reduced)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        batch_time = time.time() - end
+        end = time.time()
+        meters.update(time=batch_time)
+        if iteration % args.log_step == 0:
+            eta_seconds = meters.time.global_avg * (max_iter - iteration)
+            eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+            logger.info(
+                meters.delimiter.join([
+                    "iter: {iter:06d}",
+                    "lr: {lr:.5f}",
+                    '{meters}',
+                    "eta: {eta}",
+                    'mem: {mem}M',
+                ]).format(
+                    iter=iteration,
+                    lr=optimizer.param_groups[0]['lr'],
+                    meters=str(meters),
+                    eta=eta_string,
+                    mem=round(torch.cuda.max_memory_allocated() / 1024.0 / 1024.0),
+                )
+            )
+            if summary_writer:
+                global_step = iteration
+                summary_writer.add_scalar('losses/total_loss', losses_reduced, global_step=global_step)
+                for loss_name, loss_item in loss_dict_reduced.items():
+                    summary_writer.add_scalar('losses/{}'.format(loss_name), loss_item, global_step=global_step)
+                summary_writer.add_scalar('lr', optimizer.param_groups[0]['lr'], global_step=global_step)
+
+        if iteration % args.save_step == 0:
+            checkpointer.save("model_{:06d}".format(iteration), **arguments)
+
+        if args.eval_step > 0 and iteration % args.eval_step == 0 and not iteration == max_iter:
+            eval_results = do_evaluation(cfg, model, distributed=args.distributed, iteration=iteration)
+            if dist_util.get_rank() == 0 and summary_writer:
+                for eval_result, dataset in zip(eval_results, cfg.DATASETS.TEST):
+                    write_metric(eval_result['metrics'], 'metrics/' + dataset, summary_writer, iteration)
+            model.train()  # *IMPORTANT*: change to train mode after eval.
+
+    checkpointer.save("model_final", **arguments)
+    # compute training time
+    total_training_time = int(time.time() - start_training_time)
+    total_time_str = str(datetime.timedelta(seconds=total_training_time))
+    logger.info("Total training time: {} ({:.4f} s / it)".format(total_time_str, total_training_time / max_iter))
+    return model
+
+
+
+

ssd/layers/__init__.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from .separable_conv import SeparableConv2d
+
+__all__ = ['L2Norm', 'SeparableConv2d']
+
+
+class L2Norm(nn.Module):
+    def __init__(self, n_channels, scale):
+        super(L2Norm, self).__init__()
+        self.n_channels = n_channels
+        self.gamma = scale or None
+        self.eps = 1e-10
+        self.weight = nn.Parameter(torch.Tensor(self.n_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.constant_(self.weight, self.gamma)
+
+    def forward(self, x):
+        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
+        x = torch.div(x, norm)
+        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
+        return out

ssd/layers/separable_conv.py

@@ -0,0 +1,17 @@
+from torch import nn
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, onnx_compatible=False):
+        super().__init__()
+        ReLU = nn.ReLU if onnx_compatible else nn.ReLU6
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
+                      groups=in_channels, stride=stride, padding=padding),
+            nn.BatchNorm2d(in_channels),
+            ReLU(),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return self.conv(x)

ssd/modeling/anchors/prior_box.py

@@ -0,0 +1,54 @@
+from itertools import product
+
+import torch
+from math import sqrt
+
+
+class PriorBox:
+    def __init__(self, cfg):
+        self.image_size = cfg.INPUT.IMAGE_SIZE #300
+        prior_config = cfg.MODEL.PRIORS
+        self.feature_maps = prior_config.FEATURE_MAPS
+        self.min_sizes = prior_config.MIN_SIZES  #[30, 60, 111, 162, 213, 264]
+        self.max_sizes = prior_config.MAX_SIZES  #[60, 111, 162, 213, 264, 315]
+        self.strides = prior_config.STRIDES #[8, 16, 32, 64, 100, 300]
+        self.aspect_ratios = prior_config.ASPECT_RATIOS #[[2], [2, 3], [2, 3], [2, 3], [2], [2]]
+        self.clip = prior_config.CLIP
+
+    def __call__(self):
+        """Generate SSD Prior Boxes.
+            It returns the center, height and width of the priors. The values are relative to the image size
+            Returns:
+                priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values
+                    are relative to the image size.
+        """
+        priors = []
+        for k, f in enumerate(self.feature_maps):
+            scale = self.image_size / self.strides[k]
+            for i, j in product(range(f), repeat=2):
+                # unit center x,y
+                cx = (j + 0.5) / scale
+                cy = (i + 0.5) / scale
+
+                # small sized square box
+                size = self.min_sizes[k]
+                h = w = size / self.image_size
+                priors.append([cx, cy, w, h])
+
+                # big sized square box
+                size = sqrt(self.min_sizes[k] * self.max_sizes[k])
+                h = w = size / self.image_size
+                priors.append([cx, cy, w, h])
+
+                # change h/w ratio of the small sized box
+                size = self.min_sizes[k]
+                h = w = size / self.image_size
+                for ratio in self.aspect_ratios[k]:
+                    ratio = sqrt(ratio)
+                    priors.append([cx, cy, w * ratio, h / ratio])
+                    priors.append([cx, cy, w / ratio, h * ratio])
+
+        priors = torch.tensor(priors)
+        if self.clip:
+            priors.clamp_(max=1, min=0)       #限制priors的范围在0-1之间
+        return priors

ssd/modeling/backbone/__init__.py

@@ -0,0 +1,10 @@
+from ssd.modeling import registry
+from .vgg import VGG
+from .mobilenet import MobileNetV2
+from .efficient_net import EfficientNet
+
+__all__ = ['build_backbone', 'VGG', 'MobileNetV2', 'EfficientNet']
+
+
+def build_backbone(cfg):
+    return registry.BACKBONES[cfg.MODEL.BACKBONE.NAME](cfg, cfg.MODEL.BACKBONE.PRETRAINED)

ssd/modeling/backbone/efficient_net/__init__.py

@@ -0,0 +1,14 @@
+from ssd.modeling import registry
+from .efficient_net import EfficientNet
+
+__all__ = ['efficient_net_b3', 'EfficientNet']
+
+
+@registry.BACKBONES.register('efficient_net-b3')
+def efficient_net_b3(cfg, pretrained=True):
+    if pretrained:
+        model = EfficientNet.from_pretrained('efficientnet-b3')
+    else:
+        model = EfficientNet.from_name('efficientnet-b3')
+
+    return model

ssd/modeling/backbone/efficient_net/efficient_net.py

@@ -0,0 +1,231 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .utils import (
+    relu_fn,
+    round_filters,
+    round_repeats,
+    drop_connect,
+    Conv2dSamePadding,
+    get_model_params,
+    efficientnet_params,
+    load_pretrained_weights,
+)
+
+INDICES = {
+    'efficientnet-b3': [7, 17, 25]
+}
+
+EXTRAS = {
+    'efficientnet-b3': [
+        # in,  out, k, s, p
+        [(384, 128, 1, 1, 0), (128, 256, 3, 2, 1)],  # 5 x 5
+        [(256, 128, 1, 1, 0), (128, 256, 3, 1, 0)],  # 3 x 3
+        [(256, 128, 1, 1, 0), (128, 256, 3, 1, 0)],  # 1 x 1
+
+    ]
+}
+
+
+def add_extras(cfgs):
+    extras = nn.ModuleList()
+    for cfg in cfgs:
+        extra = []
+        for prams in cfg:
+            in_channels, out_channels, kernel_size, stride, padding = prams
+            extra.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding))
+            extra.append(nn.ReLU())
+        extras.append(nn.Sequential(*extra))
+    return extras
+
+
+class MBConvBlock(nn.Module):
+    """
+    Mobile Inverted Residual Bottleneck Block
+
+    Args:
+        block_args (namedtuple): BlockArgs, see above
+        global_params (namedtuple): GlobalParam, see above
+
+    Attributes:
+        has_se (bool): Whether the block contains a Squeeze and Excitation layer.
+    """
+
+    def __init__(self, block_args, global_params):
+        super().__init__()
+        self._block_args = block_args
+        self._bn_mom = 1 - global_params.batch_norm_momentum
+        self._bn_eps = global_params.batch_norm_epsilon
+        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
+        self.id_skip = block_args.id_skip  # skip connection and drop connect
+
+        # Expansion phase
+        inp = self._block_args.input_filters  # number of input channels
+        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
+        if self._block_args.expand_ratio != 1:
+            self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
+            self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+
+        # Depthwise convolution phase
+        k = self._block_args.kernel_size
+        s = self._block_args.stride
+        self._depthwise_conv = Conv2dSamePadding(
+            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
+            kernel_size=k, stride=s, bias=False)
+        self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+
+        # Squeeze and Excitation layer, if desired
+        if self.has_se:
+            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
+            self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
+            self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)
+
+        # Output phase
+        final_oup = self._block_args.output_filters
+        self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
+        self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)
+
+    def forward(self, inputs, drop_connect_rate=None):
+        """
+        :param inputs: input tensor
+        :param drop_connect_rate: drop connect rate (float, between 0 and 1)
+        :return: output of block
+        """
+
+        # Expansion and Depthwise Convolution
+        x = inputs
+        if self._block_args.expand_ratio != 1:
+            x = relu_fn(self._bn0(self._expand_conv(inputs)))
+        x = relu_fn(self._bn1(self._depthwise_conv(x)))
+
+        # Squeeze and Excitation
+        if self.has_se:
+            x_squeezed = F.adaptive_avg_pool2d(x, 1)
+            x_squeezed = self._se_expand(relu_fn(self._se_reduce(x_squeezed)))
+            x = torch.sigmoid(x_squeezed) * x
+
+        x = self._bn2(self._project_conv(x))
+
+        # Skip connection and drop connect
+        input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters
+        if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters:
+            if drop_connect_rate:
+                x = drop_connect(x, p=drop_connect_rate, training=self.training)
+            x = x + inputs  # skip connection
+        return x
+
+
+class EfficientNet(nn.Module):
+    """
+    An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods
+
+    Args:
+        blocks_args (list): A list of BlockArgs to construct blocks
+        global_params (namedtuple): A set of GlobalParams shared between blocks
+
+    Example:
+        model = EfficientNet.from_pretrained('efficientnet-b0')
+
+    """
+
+    def __init__(self, model_name, blocks_args=None, global_params=None):
+        super().__init__()
+        self.indices = INDICES[model_name]
+        self.extras = add_extras(EXTRAS[model_name])
+        assert isinstance(blocks_args, list), 'blocks_args should be a list'
+        assert len(blocks_args) > 0, 'block args must be greater than 0'
+        self._global_params = global_params
+        self._blocks_args = blocks_args
+
+        # Batch norm parameters
+        bn_mom = 1 - self._global_params.batch_norm_momentum
+        bn_eps = self._global_params.batch_norm_epsilon
+
+        # Stem
+        in_channels = 3  # rgb
+        out_channels = round_filters(32, self._global_params)  # number of output channels
+        self._conv_stem = Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
+        self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
+
+        # Build blocks
+        self._blocks = nn.ModuleList([])
+        for block_args in self._blocks_args:
+
+            # Update block input and output filters based on depth multiplier.
+            block_args = block_args._replace(
+                input_filters=round_filters(block_args.input_filters, self._global_params),
+                output_filters=round_filters(block_args.output_filters, self._global_params),
+                num_repeat=round_repeats(block_args.num_repeat, self._global_params)
+            )
+
+            # The first block needs to take care of stride and filter size increase.
+            self._blocks.append(MBConvBlock(block_args, self._global_params))
+            if block_args.num_repeat > 1:
+                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)
+            for _ in range(block_args.num_repeat - 1):
+                self._blocks.append(MBConvBlock(block_args, self._global_params))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for m in self.extras.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def extract_features(self, inputs):
+        """ Returns output of the final convolution layer """
+
+        # Stem
+        x = relu_fn(self._bn0(self._conv_stem(inputs)))
+
+        features = []
+
+        # Blocks
+        for idx, block in enumerate(self._blocks):
+            drop_connect_rate = self._global_params.drop_connect_rate
+            if drop_connect_rate:
+                drop_connect_rate *= float(idx) / len(self._blocks)
+            x = block(x, drop_connect_rate)
+            if idx in self.indices:
+                features.append(x)
+
+        return x, features
+
+    def forward(self, inputs):
+        """ Calls extract_features to extract features, applies final linear layer, and returns logits. """
+
+        # Convolution layers
+        x, features = self.extract_features(inputs)
+
+        for layer in self.extras:
+            x = layer(x)
+            features.append(x)
+
+        return tuple(features)
+
+    @classmethod
+    def from_name(cls, model_name, override_params=None):
+        cls._check_model_name_is_valid(model_name)
+        blocks_args, global_params = get_model_params(model_name, override_params)
+        return EfficientNet(model_name, blocks_args, global_params)
+
+    @classmethod
+    def from_pretrained(cls, model_name):
+        model = EfficientNet.from_name(model_name)
+        load_pretrained_weights(model, model_name)
+        return model
+
+    @classmethod
+    def get_image_size(cls, model_name):
+        cls._check_model_name_is_valid(model_name)
+        _, _, res, _ = efficientnet_params(model_name)
+        return res
+
+    @classmethod
+    def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False):
+        """ Validates model name. None that pretrained weights are only available for
+        the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """
+        num_models = 4 if also_need_pretrained_weights else 8
+        valid_models = ['efficientnet_b' + str(i) for i in range(num_models)]
+        if model_name.replace('-', '_') not in valid_models:
+            raise ValueError('model_name should be one of: ' + ', '.join(valid_models))

ssd/modeling/backbone/efficient_net/utils.py

@@ -0,0 +1,247 @@
+"""
+This file contains helper functions for building the model and for loading model parameters.
+These helper functions are built to mirror those in the official TensorFlow implementation.
+"""
+
+import re
+import math
+import collections
+import torch
+from torch import nn
+from torch.nn import functional as F
+from ssd.utils.model_zoo import load_state_dict_from_url
+
+########################################################################
+############### HELPERS FUNCTIONS FOR MODEL ARCHITECTURE ###############
+########################################################################
+
+
+# Parameters for the entire model (stem, all blocks, and head)
+
+GlobalParams = collections.namedtuple('GlobalParams', [
+    'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate',
+    'num_classes', 'width_coefficient', 'depth_coefficient',
+    'depth_divisor', 'min_depth', 'drop_connect_rate', ])
+
+# Parameters for an individual model block
+BlockArgs = collections.namedtuple('BlockArgs', [
+    'kernel_size', 'num_repeat', 'input_filters', 'output_filters',
+    'expand_ratio', 'id_skip', 'stride', 'se_ratio'])
+
+# Change namedtuple defaults
+GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields)
+BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields)
+
+
+def relu_fn(x):
+    """ Swish activation function """
+    return x * torch.sigmoid(x)
+
+
+def round_filters(filters, global_params):
+    """ Calculate and round number of filters based on depth multiplier. """
+    multiplier = global_params.width_coefficient
+    if not multiplier:
+        return filters
+    divisor = global_params.depth_divisor
+    min_depth = global_params.min_depth
+    filters *= multiplier
+    min_depth = min_depth or divisor
+    new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)
+    if new_filters < 0.9 * filters:  # prevent rounding by more than 10%
+        new_filters += divisor
+    return int(new_filters)
+
+
+def round_repeats(repeats, global_params):
+    """ Round number of filters based on depth multiplier. """
+    multiplier = global_params.depth_coefficient
+    if not multiplier:
+        return repeats
+    return int(math.ceil(multiplier * repeats))
+
+
+def drop_connect(inputs, p, training):
+    """ Drop connect. """
+    if not training: return inputs
+    batch_size = inputs.shape[0]
+    keep_prob = 1 - p
+    random_tensor = keep_prob
+    random_tensor += torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device)
+    binary_tensor = torch.floor(random_tensor)
+    output = inputs / keep_prob * binary_tensor
+    return output
+
+
+class Conv2dSamePadding(nn.Conv2d):
+    """ 2D Convolutions like TensorFlow """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True):
+        super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)
+        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2
+
+    def forward(self, x):
+        ih, iw = x.size()[-2:]
+        kh, kw = self.weight.size()[-2:]
+        sh, sw = self.stride
+        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
+        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
+        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+        return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+########################################################################
+############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ##############
+########################################################################
+
+
+def efficientnet_params(model_name):
+    """ Map EfficientNet model name to parameter coefficients. """
+    params_dict = {
+        # Coefficients:   width,depth,res,dropout
+        'efficientnet-b0': (1.0, 1.0, 224, 0.2),
+        'efficientnet-b1': (1.0, 1.1, 240, 0.2),
+        'efficientnet-b2': (1.1, 1.2, 260, 0.3),
+        'efficientnet-b3': (1.2, 1.4, 300, 0.3),
+        'efficientnet-b4': (1.4, 1.8, 380, 0.4),
+        'efficientnet-b5': (1.6, 2.2, 456, 0.4),
+        'efficientnet-b6': (1.8, 2.6, 528, 0.5),
+        'efficientnet-b7': (2.0, 3.1, 600, 0.5),
+    }
+    return params_dict[model_name]
+
+
+class BlockDecoder(object):
+    """ Block Decoder for readability, straight from the official TensorFlow repository """
+
+    @staticmethod
+    def _decode_block_string(block_string):
+        """ Gets a block through a string notation of arguments. """
+        assert isinstance(block_string, str)
+
+        ops = block_string.split('_')
+        options = {}
+        for op in ops:
+            splits = re.split(r'(\d.*)', op)
+            if len(splits) >= 2:
+                key, value = splits[:2]
+                options[key] = value
+
+        # Check stride
+        assert (('s' in options and len(options['s']) == 1) or
+                (len(options['s']) == 2 and options['s'][0] == options['s'][1]))
+
+        return BlockArgs(
+            kernel_size=int(options['k']),
+            num_repeat=int(options['r']),
+            input_filters=int(options['i']),
+            output_filters=int(options['o']),
+            expand_ratio=int(options['e']),
+            id_skip=('noskip' not in block_string),
+            se_ratio=float(options['se']) if 'se' in options else None,
+            stride=[int(options['s'][0])])
+
+    @staticmethod
+    def _encode_block_string(block):
+        """Encodes a block to a string."""
+        args = [
+            'r%d' % block.num_repeat,
+            'k%d' % block.kernel_size,
+            's%d%d' % (block.strides[0], block.strides[1]),
+            'e%s' % block.expand_ratio,
+            'i%d' % block.input_filters,
+            'o%d' % block.output_filters
+        ]
+        if 0 < block.se_ratio <= 1:
+            args.append('se%s' % block.se_ratio)
+        if block.id_skip is False:
+            args.append('noskip')
+        return '_'.join(args)
+
+    @staticmethod
+    def decode(string_list):
+        """
+        Decodes a list of string notations to specify blocks inside the network.
+
+        :param string_list: a list of strings, each string is a notation of block
+        :return: a list of BlockArgs namedtuples of block args
+        """
+        assert isinstance(string_list, list)
+        blocks_args = []
+        for block_string in string_list:
+            blocks_args.append(BlockDecoder._decode_block_string(block_string))
+        return blocks_args
+
+    @staticmethod
+    def encode(blocks_args):
+        """
+        Encodes a list of BlockArgs to a list of strings.
+
+        :param blocks_args: a list of BlockArgs namedtuples of block args
+        :return: a list of strings, each string is a notation of block
+        """
+        block_strings = []
+        for block in blocks_args:
+            block_strings.append(BlockDecoder._encode_block_string(block))
+        return block_strings
+
+
+def efficientnet(width_coefficient=None, depth_coefficient=None,
+                 dropout_rate=0.2, drop_connect_rate=0.2):
+    """ Creates a efficientnet model. """
+
+    blocks_args = [
+        'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25',
+        'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25',
+        'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25',
+        'r1_k3_s11_e6_i192_o320_se0.25',
+    ]
+    blocks_args = BlockDecoder.decode(blocks_args)
+
+    global_params = GlobalParams(
+        batch_norm_momentum=0.99,
+        batch_norm_epsilon=1e-3,
+        dropout_rate=dropout_rate,
+        drop_connect_rate=drop_connect_rate,
+        # data_format='channels_last',  # removed, this is always true in PyTorch
+        num_classes=1000,
+        width_coefficient=width_coefficient,
+        depth_coefficient=depth_coefficient,
+        depth_divisor=8,
+        min_depth=None
+    )
+
+    return blocks_args, global_params
+
+
+def get_model_params(model_name, override_params):
+    """ Get the block args and global params for a given model """
+    if model_name.startswith('efficientnet'):
+        w, d, _, p = efficientnet_params(model_name)
+        # note: all models have drop connect rate = 0.2
+        blocks_args, global_params = efficientnet(width_coefficient=w, depth_coefficient=d, dropout_rate=p)
+    else:
+        raise NotImplementedError('model name is not pre-defined: %s' % model_name)
+    if override_params:
+        # ValueError will be raised here if override_params has fields not included in global_params.
+        global_params = global_params._replace(**override_params)
+    return blocks_args, global_params
+
+
+url_map = {
+    'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet-b0-08094119.pth',
+    'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet-b1-dbc7070a.pth',
+    'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet-b2-27687264.pth',
+    'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet-b3-c8376fa2.pth',
+    'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet-b4-e116e8b3.pth',
+    'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet-b5-586e6cc6.pth',
+}
+
+
+def load_pretrained_weights(model, model_name):
+    """ Loads pretrained weights, and downloads if loading for the first time. """
+    state_dict = load_state_dict_from_url(url_map[model_name])
+    model.load_state_dict(state_dict, strict=False)
+    print('Loaded pretrained weights for {}'.format(model_name))

ssd/modeling/backbone/mobilenet.py

@@ -0,0 +1,134 @@
+from torch import nn
+
+from ssd.modeling import registry
+from ssd.utils.model_zoo import load_state_dict_from_url
+
+model_urls = {
+    'mobilenet_v2': 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth',
+}
+
+
+class ConvBNReLU(nn.Sequential):
+    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
+        padding = (kernel_size - 1) // 2
+        super(ConvBNReLU, self).__init__(
+            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.ReLU6(inplace=True)
+        )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(round(inp * expand_ratio))
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        layers = []
+        if expand_ratio != 1:
+            # pw
+            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
+        layers.extend([
+            # dw
+            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
+            # pw-linear
+            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(oup),
+        ])
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, width_mult=1.0, inverted_residual_setting=None):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+
+        if inverted_residual_setting is None:
+            inverted_residual_setting = [
+                # t, c, n, s
+                [1, 16, 1, 1],
+                [6, 24, 2, 2],
+                [6, 32, 3, 2],
+                [6, 64, 4, 2],
+                [6, 96, 3, 1],
+                [6, 160, 3, 2],
+                [6, 320, 1, 1],
+            ]
+
+        # only check the first element, assuming user knows t,c,n,s are required
+        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
+            raise ValueError("inverted_residual_setting should be non-empty "
+                             "or a 4-element list, got {}".format(inverted_residual_setting))
+
+        # building first layer
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * max(1.0, width_mult))
+        features = [ConvBNReLU(3, input_channel, stride=2)]
+        # building inverted residual blocks
+        for t, c, n, s in inverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                stride = s if i == 0 else 1
+                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
+                input_channel = output_channel
+        # building last several layers
+        features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
+        # make it nn.Sequential
+        self.features = nn.Sequential(*features)
+        self.extras = nn.ModuleList([
+            InvertedResidual(1280, 512, 2, 0.2),
+            InvertedResidual(512, 256, 2, 0.25),
+            InvertedResidual(256, 256, 2, 0.5),
+            InvertedResidual(256, 64, 2, 0.25)
+        ])
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # weight initialization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        features = []
+        for i in range(14):
+            x = self.features[i](x)
+        features.append(x)
+
+        for i in range(14, len(self.features)):
+            x = self.features[i](x)
+        features.append(x)
+
+        for i in range(len(self.extras)):
+            x = self.extras[i](x)
+            features.append(x)
+
+        return tuple(features)
+
+
+@registry.BACKBONES.register('mobilenet_v2')
+def mobilenet_v2(cfg, pretrained=True):
+    model = MobileNetV2()
+    if pretrained:
+        model.load_state_dict(load_state_dict_from_url(model_urls['mobilenet_v2']), strict=False)
+    return model

ssd/modeling/backbone/vgg.py

@@ -0,0 +1,114 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ssd.layers import L2Norm
+from ssd.modeling import registry
+from ssd.utils.model_zoo import load_state_dict_from_url
+
+model_urls = {
+    'vgg': 'https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth',
+}
+
+
+# borrowed from https://github.com/amdegroot/ssd.pytorch/blob/master/ssd.py
+def add_vgg(cfg, batch_norm=False):
+    layers = []
+    in_channels = 3
+    for v in cfg:
+        if v == 'M':
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+        elif v == 'C':
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
+        else:
+            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+            if batch_norm:
+                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
+    conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
+    conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
+    layers += [pool5, conv6,
+               nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
+    return layers
+
+
+def add_extras(cfg, i, size=300):
+    # Extra layers added to VGG for feature scaling
+    layers = []
+    in_channels = i
+    flag = False
+    for k, v in enumerate(cfg):
+        if in_channels != 'S':
+            if v == 'S':
+                layers += [nn.Conv2d(in_channels, cfg[k + 1], kernel_size=(1, 3)[flag], stride=2, padding=1)]
+            else:
+                layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])]
+            flag = not flag
+        in_channels = v
+    if size == 512:
+        layers.append(nn.Conv2d(in_channels, 128, kernel_size=1, stride=1))
+        layers.append(nn.Conv2d(128, 256, kernel_size=4, stride=1, padding=1))
+    return layers
+
+
+vgg_base = {
+    '300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
+            512, 512, 512],
+    '512': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
+            512, 512, 512],
+}
+extras_base = {
+    '300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256],
+    '512': [256, 'S', 512, 128, 'S', 256, 128, 'S', 256, 128, 'S', 256],
+}
+
+
+class VGG(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        size = cfg.INPUT.IMAGE_SIZE
+        vgg_config = vgg_base[str(size)]
+        extras_config = extras_base[str(size)]
+
+        self.vgg = nn.ModuleList(add_vgg(vgg_config))
+        self.extras = nn.ModuleList(add_extras(extras_config, i=1024, size=size))
+        self.l2_norm = L2Norm(512, scale=20)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for m in self.extras.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def init_from_pretrain(self, state_dict):
+        self.vgg.load_state_dict(state_dict)
+
+    def forward(self, x):
+        features = []
+        for i in range(23):
+            x = self.vgg[i](x)
+        s = self.l2_norm(x)  # Conv4_3 L2 normalization
+        features.append(s)
+
+        # apply vgg up to fc7
+        for i in range(23, len(self.vgg)):
+            x = self.vgg[i](x)
+        features.append(x)
+
+        for k, v in enumerate(self.extras):
+            x = F.relu(v(x), inplace=True)
+            if k % 2 == 1:
+                features.append(x)
+
+        return tuple(features)
+
+
+@registry.BACKBONES.register('vgg')
+def vgg(cfg, pretrained=True):
+    model = VGG(cfg)
+    if pretrained:
+        model.init_from_pretrain(load_state_dict_from_url(model_urls['vgg']))
+    return model

ssd/modeling/box_head/__init__.py

@@ -0,0 +1,8 @@
+from ssd.modeling import registry
+from .box_head import SSDBoxHead
+
+__all__ = ['build_box_head', 'SSDBoxHead']
+
+
+def build_box_head(cfg):
+    return registry.BOX_HEADS[cfg.MODEL.BOX_HEAD.NAME](cfg)

ssd/modeling/box_head/box_head.py

@@ -0,0 +1,49 @@
+from torch import nn
+import torch.nn.functional as F
+
+from ssd.modeling import registry
+from ssd.modeling.anchors.prior_box import PriorBox
+from ssd.modeling.box_head.box_predictor import make_box_predictor
+from ssd.utils import box_utils
+from .inference import PostProcessor
+from .loss import MultiBoxLoss
+
+
+@registry.BOX_HEADS.register('SSDBoxHead')
+class SSDBoxHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.predictor = make_box_predictor(cfg)
+        self.loss_evaluator = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO) #caculate loss
+        self.post_processor = PostProcessor(cfg)
+        self.priors = None
+
+    def forward(self, features, targets=None):
+        cls_logits, bbox_pred = self.predictor(features)  #返回所有卷积结果的tensor合并
+        if self.training:
+            return self._forward_train(cls_logits, bbox_pred, targets)
+        else:
+            return self._forward_test(cls_logits, bbox_pred)
+
+    def _forward_train(self, cls_logits, bbox_pred, targets):
+        gt_boxes, gt_labels = targets['boxes'], targets['labels']
+        reg_loss, cls_loss = self.loss_evaluator(cls_logits, bbox_pred, gt_labels, gt_boxes)
+        loss_dict = dict(
+            reg_loss=reg_loss,
+            cls_loss=cls_loss,
+        )
+        detections = (cls_logits, bbox_pred)
+        return detections, loss_dict
+
+    def _forward_test(self, cls_logits, bbox_pred):
+        if self.priors is None:
+            self.priors = PriorBox(self.cfg)().to(bbox_pred.device)
+        scores = F.softmax(cls_logits, dim=2)  # 计算分类损失
+        boxes = box_utils.convert_locations_to_boxes(
+            bbox_pred, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE
+        )  #将卷积的结果转化成boxex格式 bbox_pred前面卷积生成的是边框的四个坐标形式xyxy  self.priors xywh
+        boxes = box_utils.center_form_to_corner_form(boxes) #变成xyxy
+        detections = (scores, boxes)
+        detections = self.post_processor(detections)
+        return detections, {}

ssd/modeling/box_head/box_predictor.py

@@ -0,0 +1,71 @@
+import torch
+from torch import nn
+
+from ssd.layers import SeparableConv2d
+from ssd.modeling import registry
+
+
+class BoxPredictor(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.cls_headers = nn.ModuleList()
+        self.reg_headers = nn.ModuleList()
+        for level, (boxes_per_location, out_channels) in enumerate(zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)):#对要求的feature map进行分类和回归卷积层建立
+            self.cls_headers.append(self.cls_block(level, out_channels, boxes_per_location))
+            self.reg_headers.append(self.reg_block(level, out_channels, boxes_per_location))
+        self.reset_parameters()
+
+    def cls_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, features):
+        cls_logits = []
+        bbox_pred = []
+        for feature, cls_header, reg_header in zip(features, self.cls_headers, self.reg_headers):
+            cls_logits.append(cls_header(feature).permute(0, 2, 3, 1).contiguous()) #计算每一层的分类卷积结果  contiguous将结果在内存中变成连续的
+            bbox_pred.append(reg_header(feature).permute(0, 2, 3, 1).contiguous()) #计算每一层的边框回归的卷积结果
+
+        batch_size = features[0].shape[0]
+        cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1, self.cfg.MODEL.NUM_CLASSES)
+        bbox_pred = torch.cat([l.view(l.shape[0], -1) for l in bbox_pred], dim=1).view(batch_size, -1, 4)
+
+        return cls_logits, bbox_pred
+
+
+@registry.BOX_PREDICTORS.register('SSDBoxPredictor')
+class SSDBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+@registry.BOX_PREDICTORS.register('SSDLiteBoxPredictor')
+class SSDLiteBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+def make_box_predictor(cfg):
+    return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg)  #将括号中的函数当作参数进行调用
+#

ssd/modeling/box_head/inference.py

@@ -0,0 +1,65 @@
+import torch
+
+from ssd.structures.container import Container
+from ssd.utils.nms import boxes_nms
+
+
+class PostProcessor: #传进来的是边框的类别打分和边框boxes
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.width = cfg.INPUT.IMAGE_SIZE
+        self.height = cfg.INPUT.IMAGE_SIZE
+
+    def __call__(self, detections):  #detections传进来的参数 是一个三维 10*8732*21 类别 相当于每个box都对每个类别进行预测 10*8732*4 相当于对  10 是10幅图片
+        batches_scores, batches_boxes = detections
+        device = batches_scores.device
+        batch_size = batches_scores.size(0)
+        results = []
+        for batch_id in range(batch_size):  #从0到9
+            processed_boxes = []
+            processed_scores = []
+            processed_labels = []
+
+            per_img_scores, per_img_boxes = batches_scores[batch_id], batches_boxes[batch_id]  # (N, #CLS) (N, 4)取一副图片的数据
+            #print(per_img_scores.data)
+            for class_id in range(1, per_img_scores.size(1)):  # skip background
+
+                scores = per_img_scores[:, class_id]
+                mask = scores > self.cfg.TEST.CONFIDENCE_THRESHOLD
+                scores = scores[mask]
+                if scores.size(0) == 0:
+                    continue
+                boxes = per_img_boxes[mask, :]
+                boxes[:, 0::2] *= self.width
+                boxes[:, 1::2] *= self.height
+
+                keep = boxes_nms(boxes, scores, self.cfg.TEST.NMS_THRESHOLD, self.cfg.TEST.MAX_PER_CLASS)
+
+                nmsed_boxes = boxes[keep, :]
+                nmsed_labels = torch.tensor([class_id] * keep.size(0), device=device)
+                nmsed_scores = scores[keep]
+
+                processed_boxes.append(nmsed_boxes)
+                processed_scores.append(nmsed_scores)
+                processed_labels.append(nmsed_labels)
+
+            if len(processed_boxes) == 0:
+                processed_boxes = torch.empty(0, 4)
+                processed_labels = torch.empty(0)
+                processed_scores = torch.empty(0)
+            else:
+                processed_boxes = torch.cat(processed_boxes, 0)
+                processed_labels = torch.cat(processed_labels, 0)
+                processed_scores = torch.cat(processed_scores, 0)
+
+            if processed_boxes.size(0) > self.cfg.TEST.MAX_PER_IMAGE > 0:
+                processed_scores, keep = torch.topk(processed_scores, k=self.cfg.TEST.MAX_PER_IMAGE)
+                processed_boxes = processed_boxes[keep, :]
+                processed_labels = processed_labels[keep]
+
+            container = Container(boxes=processed_boxes, labels=processed_labels, scores=processed_scores)  #结果输出 处理好的边框坐标 类别 分数
+            container.img_width = self.width
+            container.img_height = self.height
+            results.append(container)
+        return results

ssd/modeling/box_head/loss.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from ssd.utils import box_utils
+
+
+class MultiBoxLoss(nn.Module):
+    def __init__(self, neg_pos_ratio):
+        """Implement SSD MultiBox Loss.
+
+        Basically, MultiBox loss combines classification loss
+         and Smooth L1 regression loss.
+        """
+        super(MultiBoxLoss, self).__init__()
+        self.neg_pos_ratio = neg_pos_ratio
+
+    def forward(self, confidence, predicted_locations, labels, gt_locations):
+        """Compute classification loss and smooth l1 loss.
+
+        Args:
+            confidence (batch_size, num_priors, num_classes): class predictions.
+            predicted_locations (batch_size, num_priors, 4): predicted locations.
+            labels (batch_size, num_priors): real labels of all the priors.
+            gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors.
+        """
+        num_classes = confidence.size(2)
+        with torch.no_grad():
+            # derived from cross_entropy=sum(log(p))
+            loss = -F.log_softmax(confidence, dim=2)[:, :, 0]
+            mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio)
+
+        confidence = confidence[mask, :]
+        classification_loss = F.cross_entropy(confidence.view(-1, num_classes), labels[mask], reduction='sum')
+
+        pos_mask = labels > 0
+        predicted_locations = predicted_locations[pos_mask, :].view(-1, 4)
+        gt_locations = gt_locations[pos_mask, :].view(-1, 4)
+        smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum')
+        num_pos = gt_locations.size(0)
+        return smooth_l1_loss / num_pos, classification_loss / num_pos

ssd/modeling/detector/__init__.py

@@ -0,0 +1,10 @@
+from .ssd_detector import SSDDetector
+
+_DETECTION_META_ARCHITECTURES = {
+    "SSDDetector": SSDDetector
+}
+
+
+def build_detection_model(cfg):
+    meta_arch = _DETECTION_META_ARCHITECTURES[cfg.MODEL.META_ARCHITECTURE]
+    return meta_arch(cfg)

ssd/modeling/detector/ssd_detector.py

@@ -0,0 +1,19 @@
+from torch import nn
+
+from ssd.modeling.backbone import build_backbone
+from ssd.modeling.box_head import build_box_head
+
+
+class SSDDetector(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.backbone = build_backbone(cfg) #vgg
+        self.box_head = build_box_head(cfg) #SSDBoxHead
+
+    def forward(self, images, targets=None):
+        features = self.backbone(images)  #get feature map
+        detections, detector_losses = self.box_head(features, targets)
+        if self.training:
+            return detector_losses
+        return detections

ssd/modeling/registry.py

@@ -0,0 +1,5 @@
+from ssd.utils.registry import Registry
+
+BACKBONES = Registry()
+BOX_HEADS = Registry()
+BOX_PREDICTORS = Registry()

ssd/solver/build.py

@@ -0,0 +1,16 @@
+import torch
+
+from .lr_scheduler import WarmupMultiStepLR
+
+
+def make_optimizer(cfg, model, lr=None):
+    lr = cfg.SOLVER.BASE_LR if lr is None else lr
+    return torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
+
+
+def make_lr_scheduler(cfg, optimizer, milestones=None):
+    return WarmupMultiStepLR(optimizer=optimizer,
+                             milestones=cfg.SOLVER.LR_STEPS if milestones is None else milestones,
+                             gamma=cfg.SOLVER.GAMMA,
+                             warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
+                             warmup_iters=cfg.SOLVER.WARMUP_ITERS)

ssd/solver/lr_scheduler.py

@@ -0,0 +1,31 @@
+from bisect import bisect_right
+
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class WarmupMultiStepLR(_LRScheduler):
+    def __init__(self, optimizer, milestones, gamma=0.1, warmup_factor=1.0 / 3,
+                 warmup_iters=500, last_epoch=-1):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                "Milestones should be a list of" " increasing integers. Got {}",
+                milestones,
+            )
+
+        self.milestones = milestones
+        self.gamma = gamma
+        self.warmup_factor = warmup_factor
+        self.warmup_iters = warmup_iters
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        warmup_factor = 1
+        if self.last_epoch < self.warmup_iters:
+            alpha = float(self.last_epoch) / self.warmup_iters
+            warmup_factor = self.warmup_factor * (1 - alpha) + alpha
+        return [
+            base_lr
+            * warmup_factor
+            * self.gamma ** bisect_right(self.milestones, self.last_epoch)
+            for base_lr in self.base_lrs
+        ]

ssd/structures/container.py

@@ -0,0 +1,54 @@
+class Container:
+    """
+    Help class for manage boxes, labels, etc...
+    Not inherit dict due to `default_collate` will change dict's subclass to dict.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._data_dict = dict(*args, **kwargs)
+
+    def __setattr__(self, key, value):
+        object.__setattr__(self, key, value)
+
+    def __getitem__(self, key):
+        return self._data_dict[key]
+
+    def __iter__(self):
+        return self._data_dict.__iter__()
+
+    def __setitem__(self, key, value):
+        self._data_dict[key] = value
+
+    def _call(self, name, *args, **kwargs):
+        keys = list(self._data_dict.keys())
+        for key in keys:
+            value = self._data_dict[key]
+            if hasattr(value, name):
+                self._data_dict[key] = getattr(value, name)(*args, **kwargs)
+        return self
+
+    def to(self, *args, **kwargs):
+        return self._call('to', *args, **kwargs)
+
+    def numpy(self):
+        return self._call('numpy')
+
+    def resize(self, size):
+        """resize boxes
+        Args:
+            size: (width, height)
+        Returns:
+            self
+        """
+        img_width = getattr(self, 'img_width', -1)
+        img_height = getattr(self, 'img_height', -1)
+        assert img_width > 0 and img_height > 0
+        assert 'boxes' in self._data_dict
+        boxes = self._data_dict['boxes']
+        new_width, new_height = size
+        boxes[:, 0::2] *= (new_width / img_width)
+        boxes[:, 1::2] *= (new_height / img_height)
+        return self
+
+    def __repr__(self):
+        return self._data_dict.__repr__()

ssd/utils/__init__.py

@@ -0,0 +1 @@
+from .misc import *

ssd/utils/box_utils.py

@@ -0,0 +1,138 @@
+import torch
+import math
+
+
+def convert_locations_to_boxes(locations, priors, center_variance,
+                               size_variance):
+    """Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w).
+
+    The conversion:
+        $$predicted\_center * center_variance = \frac {real\_center - prior\_center} {prior\_hw}$$
+        $$exp(predicted\_hw * size_variance) = \frac {real\_hw} {prior\_hw}$$
+    We do it in the inverse direction here.
+    Args:
+        locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well.
+        priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes.
+        center_variance: a float used to change the scale of center.
+        size_variance: a float used to change of scale of size.
+    Returns:
+        boxes:  priors: [[center_x, center_y, w, h]]. All the values
+            are relative to the image size.
+    """
+    # priors can have one dimension less.
+    if priors.dim() + 1 == locations.dim():
+        priors = priors.unsqueeze(0)
+    return torch.cat([
+        locations[..., :2] * center_variance * priors[..., 2:] + priors[..., :2],
+        torch.exp(locations[..., 2:] * size_variance) * priors[..., 2:]
+    ], dim=locations.dim() - 1)
+
+
+def convert_boxes_to_locations(center_form_boxes, center_form_priors, center_variance, size_variance):
+    # priors can have one dimension less
+    if center_form_priors.dim() + 1 == center_form_boxes.dim():
+        center_form_priors = center_form_priors.unsqueeze(0)
+    return torch.cat([
+        (center_form_boxes[..., :2] - center_form_priors[..., :2]) / center_form_priors[..., 2:] / center_variance,
+        torch.log(center_form_boxes[..., 2:] / center_form_priors[..., 2:]) / size_variance
+    ], dim=center_form_boxes.dim() - 1)
+
+
+def area_of(left_top, right_bottom) -> torch.Tensor:
+    """Compute the areas of rectangles given two corners.
+
+    Args:
+        left_top (N, 2): left top corner.
+        right_bottom (N, 2): right bottom corner.
+
+    Returns:
+        area (N): return the area.
+    """
+    hw = torch.clamp(right_bottom - left_top, min=0.0)
+    return hw[..., 0] * hw[..., 1]
+
+
+def iou_of(boxes0, boxes1, eps=1e-5):
+    """Return intersection-over-union (Jaccard index) of boxes.
+
+    Args:
+        boxes0 (N, 4): ground truth boxes.
+        boxes1 (N or 1, 4): predicted boxes.
+        eps: a small number to avoid 0 as denominator.
+    Returns:
+        iou (N): IoU values.
+    """
+    overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2])
+    overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:])
+
+    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
+    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
+    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
+    return overlap_area / (area0 + area1 - overlap_area + eps)
+
+
+def assign_priors(gt_boxes, gt_labels, corner_form_priors,
+                  iou_threshold):
+    """Assign ground truth boxes and targets to priors.
+
+    Args:
+        gt_boxes (num_targets, 4): ground truth boxes.
+        gt_labels (num_targets): labels of targets.
+        priors (num_priors, 4): corner form priors
+    Returns:
+        boxes (num_priors, 4): real values for priors.
+        labels (num_priros): labels for priors.
+    """
+    # size: num_priors x num_targets
+    ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1))
+    # size: num_priors
+    best_target_per_prior, best_target_per_prior_index = ious.max(1)
+    # size: num_targets
+    best_prior_per_target, best_prior_per_target_index = ious.max(0)
+
+    for target_index, prior_index in enumerate(best_prior_per_target_index):
+        best_target_per_prior_index[prior_index] = target_index
+    # 2.0 is used to make sure every target has a prior assigned
+    best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2)
+    # size: num_priors
+    labels = gt_labels[best_target_per_prior_index]
+    labels[best_target_per_prior < iou_threshold] = 0  # the backgournd id
+    boxes = gt_boxes[best_target_per_prior_index]
+    return boxes, labels
+
+
+def hard_negative_mining(loss, labels, neg_pos_ratio):
+    """
+    It used to suppress the presence of a large number of negative prediction.
+    It works on image level not batch level.
+    For any example/image, it keeps all the positive predictions and
+     cut the number of negative predictions to make sure the ratio
+     between the negative examples and positive examples is no more
+     the given ratio for an image.
+
+    Args:
+        loss (N, num_priors): the loss for each example.
+        labels (N, num_priors): the labels.
+        neg_pos_ratio:  the ratio between the negative examples and positive examples.
+    """
+    pos_mask = labels > 0
+    num_pos = pos_mask.long().sum(dim=1, keepdim=True)
+    num_neg = num_pos * neg_pos_ratio
+
+    loss[pos_mask] = -math.inf
+    _, indexes = loss.sort(dim=1, descending=True)
+    _, orders = indexes.sort(dim=1)
+    neg_mask = orders < num_neg
+    return pos_mask | neg_mask
+
+
+def center_form_to_corner_form(locations):
+    return torch.cat([locations[..., :2] - locations[..., 2:] / 2,
+                      locations[..., :2] + locations[..., 2:] / 2], locations.dim() - 1)
+
+
+def corner_form_to_center_form(boxes):
+    return torch.cat([
+        (boxes[..., :2] + boxes[..., 2:]) / 2,
+        boxes[..., 2:] - boxes[..., :2]
+    ], boxes.dim() - 1)

ssd/utils/checkpoint.py

@@ -0,0 +1,107 @@
+import logging
+import os
+
+import torch
+from torch.nn.parallel import DistributedDataParallel
+
+from ssd.utils.model_zoo import cache_url
+
+
+class CheckPointer:
+    _last_checkpoint_name = 'last_checkpoint.txt'
+
+    def __init__(self,
+                 model,
+                 optimizer=None,
+                 scheduler=None,
+                 save_dir="",
+                 save_to_disk=None,
+                 logger=None):
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.save_dir = save_dir
+        self.save_to_disk = save_to_disk
+        if logger is None:
+            logger = logging.getLogger(__name__)
+        self.logger = logger
+
+    def save(self, name, **kwargs):
+        if not self.save_dir:
+            return
+
+        if not self.save_to_disk:
+            return
+
+        data = {}
+        if isinstance(self.model, DistributedDataParallel):
+            data['model'] = self.model.module.state_dict()
+        else:
+            data['model'] = self.model.state_dict()
+        if self.optimizer is not None:
+            data["optimizer"] = self.optimizer.state_dict()
+        if self.scheduler is not None:
+            data["scheduler"] = self.scheduler.state_dict()
+        data.update(kwargs)
+
+        save_file = os.path.join(self.save_dir, "{}.pth".format(name))
+        self.logger.info("Saving checkpoint to {}".format(save_file))
+        torch.save(data, save_file)
+
+        self.tag_last_checkpoint(save_file)
+
+    def load(self, f=None, use_latest=True):
+        if self.has_checkpoint() and use_latest:
+            # override argument with existing checkpoint
+            f = self.get_checkpoint_file()
+        if not f:
+            # no checkpoint could be found
+            self.logger.info("No checkpoint found.")
+            return {}
+
+        self.logger.info("Loading checkpoint from {}".format(f))
+        checkpoint = self._load_file(f)
+        model = self.model
+        if isinstance(model, DistributedDataParallel):
+            model = self.model.module
+
+        model.load_state_dict(checkpoint.pop("model"))
+        if "optimizer" in checkpoint and self.optimizer:
+            self.logger.info("Loading optimizer from {}".format(f))
+            self.optimizer.load_state_dict(checkpoint.pop("optimizer"))
+        if "scheduler" in checkpoint and self.scheduler:
+            self.logger.info("Loading scheduler from {}".format(f))
+            self.scheduler.load_state_dict(checkpoint.pop("scheduler"))
+
+        # return any further checkpoint data
+        return checkpoint
+
+    def get_checkpoint_file(self):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        try:
+            with open(save_file, "r") as f:
+                last_saved = f.read()
+                last_saved = last_saved.strip()
+        except IOError:
+            # if file doesn't exist, maybe because it has just been
+            # deleted by a separate process
+            last_saved = ""
+        return last_saved
+
+    def has_checkpoint(self):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        return os.path.exists(save_file)
+
+    def tag_last_checkpoint(self, last_filename):
+        save_file = os.path.join(self.save_dir, self._last_checkpoint_name)
+        with open(save_file, "w") as f:
+            f.write(last_filename)
+
+    def _load_file(self, f):
+        # download url files
+        if f.startswith("http"):
+            # if the file is a url path, download it and cache it
+            cached_f = cache_url(f)
+            self.logger.info("url {} cached in {}".format(f, cached_f))
+            f = cached_f
+        return torch.load(f, map_location=torch.device("cpu"))

ssd/utils/dist_util.py

@@ -0,0 +1,96 @@
+import pickle
+
+import torch
+import torch.distributed as dist
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def synchronize():
+    """
+       Helper function to synchronize (barrier) among all processes when
+       using distributed training
+    """
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        return
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return
+    dist.barrier()
+
+
+def _encode(encoded_data, data):
+    # gets a byte representation for the data
+    encoded_bytes = pickle.dumps(data)
+    # convert this byte string into a byte tensor
+    storage = torch.ByteStorage.from_buffer(encoded_bytes)
+    tensor = torch.ByteTensor(storage).to("cuda")
+    # encoding: first byte is the size and then rest is the data
+    s = tensor.numel()
+    assert s <= 255, "Can't encode data greater than 255 bytes"
+    # put the encoded data in encoded_data
+    encoded_data[0] = s
+    encoded_data[1: (s + 1)] = tensor
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.LongTensor([tensor.numel()]).to("cuda")
+    size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list

ssd/utils/logger.py

@@ -0,0 +1,22 @@
+import logging
+import os
+import sys
+
+
+def setup_logger(name, distributed_rank, save_dir=None):
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    # don't log results for the non-master process
+    if distributed_rank > 0:
+        return logger
+    stream_handler = logging.StreamHandler(stream=sys.stdout)
+    stream_handler.setLevel(logging.DEBUG)
+    formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+    if save_dir:
+        fh = logging.FileHandler(os.path.join(save_dir, 'log.txt'))
+        fh.setLevel(logging.DEBUG)
+        fh.setFormatter(formatter)
+        logger.addHandler(fh)
+    return logger

ssd/utils/metric_logger.py

@@ -0,0 +1,66 @@
+from collections import deque, defaultdict
+import numpy as np
+import torch
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=10):
+        self.deque = deque(maxlen=window_size)
+        self.value = np.nan
+        self.series = []
+        self.total = 0.0
+        self.count = 0
+
+    def update(self, value):
+        self.deque.append(value)
+        self.series.append(value)
+        self.count += 1
+        self.total += value
+        self.value = value
+
+    @property
+    def median(self):
+        values = np.array(self.deque)
+        return np.median(values)
+
+    @property
+    def avg(self):
+        values = np.array(self.deque)
+        return np.mean(values)
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+
+class MetricLogger:
+    def __init__(self, delimiter=", "):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {:.3f} ({:.3f})".format(name, meter.avg, meter.global_avg)
+            )
+        return self.delimiter.join(loss_str)

ssd/utils/misc.py

@@ -0,0 +1,14 @@
+import errno
+import os
+
+
+def str2bool(s):
+    return s.lower() in ('true', '1')
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise

ssd/utils/model_zoo.py

@@ -0,0 +1,67 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+import os
+import sys
+
+import torch
+
+from ssd.utils.dist_util import is_main_process, synchronize
+
+try:
+    from torch.hub import _download_url_to_file
+    from torch.hub import urlparse
+    from torch.hub import HASH_REGEX
+except ImportError:
+    from torch.utils.model_zoo import _download_url_to_file
+    from torch.utils.model_zoo import urlparse
+    from torch.utils.model_zoo import HASH_REGEX
+
+
+# very similar to https://github.com/pytorch/pytorch/blob/master/torch/utils/model_zoo.py
+# but with a few improvements and modifications
+def cache_url(url, model_dir=None, progress=True):
+    r"""Loads the Torch serialized object at the given URL.
+    If the object is already present in `model_dir`, it's deserialized and
+    returned. The filename part of the URL should follow the naming convention
+    ``filename-<sha256>.ext`` where ``<sha256>`` is the first eight or more
+    digits of the SHA256 hash of the contents of the file. The hash is used to
+    ensure unique names and to verify the contents of the file.
+    The default value of `model_dir` is ``$TORCH_HOME/models`` where
+    ``$TORCH_HOME`` defaults to ``~/.torch``. The default directory can be
+    overridden with the ``$TORCH_MODEL_ZOO`` environment variable.
+    Args:
+        url (string): URL of the object to download
+        model_dir (string, optional): directory in which to save the object
+        progress (bool, optional): whether or not to display a progress bar to stderr
+    Example:
+        >>> cached_file = maskrcnn_benchmark.utils.model_zoo.cache_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth')
+    """
+    if model_dir is None:
+        torch_home = os.path.expanduser(os.getenv("TORCH_HOME", "~/.torch"))
+        model_dir = os.getenv("TORCH_MODEL_ZOO", os.path.join(torch_home, "models"))
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if filename == "model_final.pkl":
+        # workaround as pre-trained Caffe2 models from Detectron have all the same filename
+        # so make the full path the filename by replacing / with _
+        filename = parts.path.replace("/", "_")
+    cached_file = os.path.join(model_dir, filename)
+    if not os.path.exists(cached_file) and is_main_process():
+        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
+        hash_prefix = HASH_REGEX.search(filename)
+        if hash_prefix is not None:
+            hash_prefix = hash_prefix.group(1)
+            # workaround: Caffe2 models don't have a hash, but follow the R-50 convention,
+            # which matches the hash PyTorch uses. So we skip the hash matching
+            # if the hash_prefix is less than 6 characters
+            if len(hash_prefix) < 6:
+                hash_prefix = None
+        _download_url_to_file(url, cached_file, hash_prefix, progress=progress)
+    synchronize()
+    return cached_file
+
+
+def load_state_dict_from_url(url, map_location='cpu'):
+    cached_file = cache_url(url)
+    return torch.load(cached_file, map_location=map_location)

ssd/utils/nms.py

@@ -0,0 +1,33 @@
+import warnings
+import torchvision
+
+try:
+    import torch_extension
+
+    _nms = torch_extension.nms
+except ImportError:
+    if torchvision.__version__ >= '0.3.0':
+        _nms = torchvision.ops.nms
+    else:
+        from .python_nms import python_nms
+
+        _nms = python_nms
+        warnings.warn('You are using python version NMS, which is very very slow. Try compile c++ NMS '
+                      'using `cd ext & python build.py build_ext develop`')
+
+
+def boxes_nms(boxes, scores, nms_thresh, max_count=-1):
+    """ Performs non-maximum suppression, run on GPU or CPU according to
+    boxes's device.
+    Args:
+        boxes(Tensor): `xyxy` mode boxes, use absolute coordinates(or relative coordinates), shape is (n, 4)
+        scores(Tensor): scores, shape is (n, )
+        nms_thresh(float): thresh
+        max_count (int): if > 0, then only the top max_proposals are kept  after non-maximum suppression
+    Returns:
+        indices kept.
+    """
+    keep = _nms(boxes, scores, nms_thresh)
+    if max_count > 0:
+        keep = keep[:max_count]
+    return keep

ssd/utils/python_nms.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+
+
+def python_nms(boxes, scores, nms_thresh):
+    """ Performs non-maximum suppression using numpy
+        Args:
+            boxes(Tensor): `xyxy` mode boxes, use absolute coordinates(not support relative coordinates),
+                shape is (n, 4)
+            scores(Tensor): scores, shape is (n, )
+            nms_thresh(float): thresh
+        Returns:
+            indices kept.
+    """
+    if boxes.numel() == 0:
+        return torch.empty((0,), dtype=torch.long)
+    # Use numpy to run nms. Running nms in PyTorch code on CPU is really slow.
+    origin_device = boxes.device
+    cpu_device = torch.device('cpu')
+    boxes = boxes.to(cpu_device).numpy()
+    scores = scores.to(cpu_device).numpy()
+
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    areas = (x2 - x1) * (y2 - y1)
+    order = np.argsort(scores)[::-1]
+    num_detections = boxes.shape[0]
+    suppressed = np.zeros((num_detections,), dtype=np.bool)
+    for _i in range(num_detections):
+        i = order[_i]
+        if suppressed[i]:
+            continue
+        ix1 = x1[i]
+        iy1 = y1[i]
+        ix2 = x2[i]
+        iy2 = y2[i]
+        iarea = areas[i]
+
+        for _j in range(_i + 1, num_detections):
+            j = order[_j]
+            if suppressed[j]:
+                continue
+
+            xx1 = max(ix1, x1[j])
+            yy1 = max(iy1, y1[j])
+            xx2 = min(ix2, x2[j])
+            yy2 = min(iy2, y2[j])
+            w = max(0, xx2 - xx1)
+            h = max(0, yy2 - yy1)
+
+            inter = w * h
+            ovr = inter / (iarea + areas[j] - inter)
+            if ovr >= nms_thresh:
+                suppressed[j] = True
+    keep = np.nonzero(suppressed == 0)[0]
+    keep = torch.from_numpy(keep).to(origin_device)
+    return keep

ssd/utils/registry.py

@@ -0,0 +1,40 @@
+def _register_generic(module_dict, module_name, module):
+    assert module_name not in module_dict
+    module_dict[module_name] = module
+
+
+class Registry(dict):
+    """
+    A helper class for managing registering modules, it extends a dictionary
+    and provides a register functions.
+    Eg. creating a registry:
+        some_registry = Registry({"default": default_module})
+    There're two ways of registering new modules:
+    1): normal way is just calling register function:
+        def foo():
+            ...
+        some_registry.register("foo_module", foo)
+    2): used as decorator when declaring the module:
+        @some_registry.register("foo_module")
+        @some_registry.register("foo_module_nickname")
+        def foo():
+            ...
+    Access of module is just like using a dictionary, eg:
+        f = some_registry["foo_module"]
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(Registry, self).__init__(*args, **kwargs)
+
+    def register(self, module_name, module=None):
+        # used as function call
+        if module is not None:
+            _register_generic(self, module_name, module)
+            return
+
+        # used as decorator
+        def register_fn(fn):
+            _register_generic(self, module_name, fn)
+            return fn
+
+        return register_fn

ssd/utils/timer.py

@@ -0,0 +1,46 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+
+import time
+import datetime
+
+
+class Timer(object):
+    def __init__(self):
+        self.reset()
+
+    @property
+    def average_time(self):
+        return self.total_time / self.calls if self.calls > 0 else 0.0
+
+    def tic(self):
+        # using time.time instead of time.clock because time time.clock
+        # does not normalize for multithreading
+        self.start_time = time.time()
+
+    def toc(self, average=True):
+        self.add(time.time() - self.start_time)
+        if average:
+            return self.average_time
+        else:
+            return self.diff
+
+    def add(self, time_diff):
+        self.diff = time_diff
+        self.total_time += self.diff
+        self.calls += 1
+
+    def reset(self):
+        self.total_time = 0.0
+        self.calls = 0
+        self.start_time = 0.0
+        self.diff = 0.0
+
+    def avg_time_str(self):
+        time_str = str(datetime.timedelta(seconds=self.average_time))
+        return time_str
+
+
+def get_time_str(time_diff):
+    time_str = str(datetime.timedelta(seconds=time_diff))
+    return time_str