first commit

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