[MaskRCNN] move dataset loader and evaluation together

examples/FasterRCNN/common.py

@@ -95,12 +95,11 @@ def segmentation_to_mask(polys, height, width):
     Args:
         polys: a list of nx2 float array. Each array contains many (x, y) coordinates.
 
-
     Returns:
         a binary matrix of (height, width)
     """
-    polys = np.asarray([p.flatten() for p in polys], dtype='float32')
-    assert polys.size > 0, "Polygons are empty!"
+    polys = [p.flatten().tolist() for p in polys]
+    assert len(polys) > 0, "Polygons are empty!"
 
     import pycocotools.mask as cocomask
     rles = cocomask.frPyObjects(polys, height, width)

examples/FasterRCNN/config.py

@@ -85,7 +85,7 @@ _C.DATA.BASEDIR = '/path/to/your/COCO/DIR'
 _C.DATA.TRAIN = ['train2014', 'valminusminival2014']   # i.e. trainval35k, AKA train2017
 # Each VAL dataset will be evaluated separately (instead of concatenated)
 _C.DATA.VAL = ('minival2014', )  # AKA val2017
-_C.DATA.NUM_CATEGORY = 80    # 80 categories in COCO
+_C.DATA.NUM_CATEGORY = 0  # without the background class (e.g., 80 for COCO)
 _C.DATA.CLASS_NAMES = []  # NUM_CLASS (NUM_CATEGORY+1) strings, the first is "BG".
 # For COCO, this list will be populated later by the COCO data loader.
 
@@ -126,7 +126,7 @@ _C.TRAIN.LR_SCHEDULE = [240000, 320000, 360000]      # "2x" schedule in detectro
 # Longer schedules for from-scratch training (https://arxiv.org/abs/1811.08883):
 # _C.TRAIN.LR_SCHEDULE = [960000, 1040000, 1080000]    # "6x" schedule in detectron
 # _C.TRAIN.LR_SCHEDULE = [1500000, 1580000, 1620000]   # "9x" schedule in detectron
-_C.TRAIN.EVAL_PERIOD = 25  # period (epochs) to run eva
+_C.TRAIN.EVAL_PERIOD = 25  # period (epochs) to run evaluation
 
 # preprocessing --------------------
 # Alternative old (worse & faster) setting: 600
@@ -241,7 +241,7 @@ def finalize_configs(is_training):
     if is_training:
         train_scales = _C.PREPROC.TRAIN_SHORT_EDGE_SIZE
         if isinstance(train_scales, (list, tuple)) and train_scales[1] - train_scales[0] > 100:
-            # don't warmup if augmentation is on
+            # don't autotune if augmentation is on
             os.environ['TF_CUDNN_USE_AUTOTUNE'] = '0'
         os.environ['TF_AUTOTUNE_THRESHOLD'] = '1'
         assert _C.TRAINER in ['horovod', 'replicated'], _C.TRAINER

examples/FasterRCNN/data.py

@@ -10,10 +10,10 @@ from tensorpack.dataflow import (
 from tensorpack.utils import logger
 from tensorpack.utils.argtools import log_once, memoized
 
-from coco import COCODetection
 from common import (
     CustomResize, DataFromListOfDict, box_to_point8, filter_boxes_inside_shape, point8_to_box, segmentation_to_mask)
 from config import config as cfg
+from coco import DetectionDataset
 from utils.generate_anchors import generate_anchors
 from utils.np_box_ops import area as np_area
 from utils.np_box_ops import ioa as np_ioa
@@ -280,25 +280,7 @@ def get_train_dataflow():
     If MODE_MASK, gt_masks: (N, h, w)
     """
 
-    roidbs = COCODetection.load_many(
-        cfg.DATA.BASEDIR, cfg.DATA.TRAIN, add_gt=True, add_mask=cfg.MODE_MASK)
-    """
-    To train on your own data, change this to your loader.
-    Produce "roidbs" as a list of dict, in the dict the following keys are needed for training:
-    height, width: integer
-    file_name: str, full path to the image
-    boxes: numpy array of kx4 floats
-    class: numpy array of k integers
-    is_crowd: k booleans. Use k False if you don't know what it means.
-    segmentation: k lists of numpy arrays (one for each box).
-        Each list of numpy arrays corresponds to the mask for one instance.
-        Each numpy array in the list is a polygon of shape Nx2,
-        because one mask can be represented by N polygons.
-
-        If your segmentation annotations are originally masks rather than polygons,
-        either convert it, or the augmentation code below will need to be
-        changed or skipped accordingly.
-    """
+    roidbs = DetectionDataset().load_training_roidbs(cfg.DATA.TRAIN)
 
     # Valid training images should have at least one fg box.
     # But this filter shall not be applied for testing.
@@ -387,13 +369,7 @@ def get_eval_dataflow(name, shard=0, num_shards=1):
         name (str): name of the dataset to evaluate
         shard, num_shards: to get subset of evaluation data
     """
-    roidbs = COCODetection.load_many(cfg.DATA.BASEDIR, name, add_gt=False)
-    """
-    To inference on your own data, change this to your loader.
-    Produce "roidbs" as a list of dict, in the dict the following keys are needed for training:
-    file_name: str, full path to the image
-    id: an id of this image
-    """
+    roidbs = DetectionDataset().load_inference_roidbs(name)
 
     num_imgs = len(roidbs)
     img_per_shard = num_imgs // num_shards

examples/FasterRCNN/eval.py

@@ -3,19 +3,15 @@
 
 import itertools
 import numpy as np
-import os
 from collections import namedtuple
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 import cv2
 import pycocotools.mask as cocomask
 import tqdm
-from pycocotools.coco import COCO
-from pycocotools.cocoeval import COCOeval
 
 from tensorpack.utils.utils import get_tqdm_kwargs
 
-from coco import COCOMeta
 from common import CustomResize, clip_boxes
 from config import config as cfg
 
@@ -103,7 +99,8 @@ def eval_coco(df, detect_func, tqdm_bar=None):
             will create a new one.
 
     Returns:
-        list of dict, to be dumped to COCO json format
+        list of dict, in the format used by
+        `DetectionDataset.eval_or_save_inference_results`
     """
     df.reset_state()
     all_results = []
@@ -115,15 +112,10 @@ def eval_coco(df, detect_func, tqdm_bar=None):
         for img, img_id in df:
             results = detect_func(img)
             for r in results:
-                box = r.box
-                cat_id = COCOMeta.class_id_to_category_id[r.class_id]
-                box[2] -= box[0]
-                box[3] -= box[1]
-
                 res = {
                     'image_id': img_id,
-                    'category_id': cat_id,
-                    'bbox': list(map(lambda x: round(float(x), 3), box)),
+                    'category_id': r.class_id,
+                    'bbox': list(r.box),
                     'score': round(float(r.score), 4),
                 }
 
@@ -147,7 +139,8 @@ def multithread_eval_coco(dataflows, detect_funcs):
         detect_funcs: a list of callable to be used in :func:`eval_coco`
 
     Returns:
-        list of dict, to be dumped to COCO json format
+        list of dict, in the format used by
+        `DetectionDataset.eval_or_save_inference_results`
     """
     num_worker = len(dataflows)
     assert len(dataflows) == len(detect_funcs)
@@ -158,37 +151,3 @@ def multithread_eval_coco(dataflows, detect_funcs):
             futures.append(executor.submit(eval_coco, dataflow, pred, pbar))
         all_results = list(itertools.chain(*[fut.result() for fut in futures]))
         return all_results
-
-
-# https://github.com/pdollar/coco/blob/master/PythonAPI/pycocoEvalDemo.ipynb
-def print_coco_metrics(dataset, json_file):
-    """
-    Args:
-        dataset (str): name of the dataset
-        json_file (str): path to the results json file in coco format
-
-    If your data is not in COCO format, write your own evaluation function.
-    """
-    ret = {}
-    assert cfg.DATA.BASEDIR and os.path.isdir(cfg.DATA.BASEDIR)
-    annofile = os.path.join(
-        cfg.DATA.BASEDIR, 'annotations',
-        'instances_{}.json'.format(dataset))
-    coco = COCO(annofile)
-    cocoDt = coco.loadRes(json_file)
-    cocoEval = COCOeval(coco, cocoDt, 'bbox')
-    cocoEval.evaluate()
-    cocoEval.accumulate()
-    cocoEval.summarize()
-    fields = ['IoU=0.5:0.95', 'IoU=0.5', 'IoU=0.75', 'small', 'medium', 'large']
-    for k in range(6):
-        ret['mAP(bbox)/' + fields[k]] = cocoEval.stats[k]
-
-    if cfg.MODE_MASK:
-        cocoEval = COCOeval(coco, cocoDt, 'segm')
-        cocoEval.evaluate()
-        cocoEval.accumulate()
-        cocoEval.summarize()
-        for k in range(6):
-            ret['mAP(segm)/' + fields[k]] = cocoEval.stats[k]
-    return ret

examples/FasterRCNN/train.py

@@ -22,11 +22,11 @@ from tensorpack.tfutils.summary import add_moving_summary
 import model_frcnn
 import model_mrcnn
 from basemodel import image_preprocess, resnet_c4_backbone, resnet_conv5, resnet_fpn_backbone
-from coco import COCODetection
+from coco import DetectionDataset
 from config import config as cfg
 from config import finalize_configs
 from data import get_all_anchors, get_all_anchors_fpn, get_eval_dataflow, get_train_dataflow
-from eval import DetectionResult, detect_one_image, eval_coco, multithread_eval_coco, print_coco_metrics
+from eval import DetectionResult, detect_one_image, eval_coco, multithread_eval_coco
 from model_box import RPNAnchors, clip_boxes, crop_and_resize, roi_align
 from model_cascade import CascadeRCNNHead
 from model_fpn import fpn_model, generate_fpn_proposals, multilevel_roi_align, multilevel_rpn_losses
@@ -388,15 +388,13 @@ def offline_evaluate(pred_config, output_file):
         logger.info("Evaluating {} ...".format(dataset))
         dataflows = [
             get_eval_dataflow(dataset, shard=k, num_shards=num_gpu)
-            for k in range(num_gpu) ]
+            for k in range(num_gpu)]
         if num_gpu > 1:
             all_results = multithread_eval_coco(dataflows, predictors)
         else:
             all_results = eval_coco(dataflows[0], predictors[0])
         output = output_file + '-' + dataset
-        with open(output, 'w') as f:
-            json.dump(all_results, f)
-        print_coco_metrics(dataset, output)
+        DetectionDataset().eval_or_save_inference_results(all_results, dataset, output)
 
 
 def predict(pred_func, input_file):
@@ -484,14 +482,11 @@ class EvalCallback(Callback):
 
         output_file = os.path.join(
             logdir, '{}-outputs{}.json'.format(self._eval_dataset, self.global_step))
-        with open(output_file, 'w') as f:
-            json.dump(all_results, f)
-        try:
-            scores = print_coco_metrics(self._eval_dataset, output_file)
-            for k, v in scores.items():
-                self.trainer.monitors.put_scalar(k, v)
-        except Exception:
-            logger.exception("Exception in COCO evaluation.")
+
+        scores = DetectionDataset().eval_or_save_inference_results(
+            all_results, self._eval_dataset, output_file)
+        for k, v in scores.items():
+            self.trainer.monitors.put_scalar(k, v)
 
     def _trigger_epoch(self):
         if self.epoch_num in self.epochs_to_eval:
@@ -520,6 +515,7 @@ if __name__ == '__main__':
         cfg.update_args(args.config)
 
     MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model()
+    DetectionDataset()  # initialize the config with information from our dataset
 
     if args.visualize or args.evaluate or args.predict:
         assert tf.test.is_gpu_available()
@@ -538,7 +534,6 @@ if __name__ == '__main__':
                 input_names=MODEL.get_inference_tensor_names()[0],
                 output_names=MODEL.get_inference_tensor_names()[1])
             if args.predict:
-                COCODetection(cfg.DATA.BASEDIR, 'val2014')   # Only to load the class names into caches
                 predict(OfflinePredictor(predcfg), args.predict)
             elif args.evaluate:
                 assert args.evaluate.endswith('.json'), args.evaluate
@@ -573,7 +568,6 @@ if __name__ == '__main__':
         total_passes = cfg.TRAIN.LR_SCHEDULE[-1] * 8 / train_dataflow.size()
         logger.info("Total passes of the training set is: {:.5g}".format(total_passes))
 
-
         callbacks = [
             PeriodicCallback(
                 ModelSaver(max_to_keep=10, keep_checkpoint_every_n_hours=1),

tensorpack/tfutils/sessinit.py

@@ -214,7 +214,7 @@ class DictRestore(SessionInit):
         mismatch.log()
 
         upd = SessionUpdate(sess, [v for v in variables if v.name in intersect])
-        logger.info("Restoring from dict ...")
+        logger.info("Restoring {} variables from dict ...".format(len(intersect)))
         upd.update({name: value for name, value in six.iteritems(self._prms) if name in intersect})