Maskrcnn-benchmark

diff --git a/maskrcnn_benchmark/config/paths_catalog.py b/maskrcnn_benchmark/config/paths_catalog.py
index 522ef88..f325888 100644
--- a/maskrcnn_benchmark/config/paths_catalog.py
+++ b/maskrcnn_benchmark/config/paths_catalog.py
@@ -7,6 +7,18 @@ import os
 class DatasetCatalog(object):
     DATA_DIR = "datasets"
     DATASETS = {
+        "companyName_190207_train": {
+            "img_dir": "companyName190207/train",
+            "ann_file": "companyName190207/annotations/train.json"
+        },
+        "companyName_190207_val": {
+            "img_dir": "companyName190207/val",
+            "ann_file": "companyName190207/annotations/val.json"
+        },
+        "companyName_190207_test": {
+            "img_dir": "companyName190207/test",
+            "ann_file": "companyName190207/annotations/test.json"
+        },
         "coco_2017_train": {
             "img_dir": "coco/train2017",
             "ann_file": "coco/annotations/instances_train2017.json"
@@ -92,7 +104,18 @@ class DatasetCatalog(object):

     @staticmethod
     def get(name):
-        if "coco" in name:
+        if "companyName" in name:
+            data_dir = DatasetCatalog.DATA_DIR
+            attrs = DatasetCatalog.DATASETS[name]
+            args = dict(
+                root=os.path.join(data_dir, attrs["img_dir"]),
+                ann_file=os.path.join(data_dir, attrs["ann_file"]),
+            )
+            return dict(
+                factory="COCODataset",
+                args=args,
+            )
+        elif "coco" in name:
             data_dir = DatasetCatalog.DATA_DIR
             attrs = DatasetCatalog.DATASETS[name]
             args = dict(

diff --git a/maskrcnn_benchmark/data/datasets/coco.py b/maskrcnn_benchmark/data/datasets/coco.py
index 4d74ab2..46bf6bb 100644
--- a/maskrcnn_benchmark/data/datasets/coco.py
+++ b/maskrcnn_benchmark/data/datasets/coco.py
@@ -6,6 +6,13 @@ from maskrcnn_benchmark.structures.bounding_box import BoxList
 from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask


+def iscrowd_is_0(obj):
+    try:
+        return obj["iscrowd"] == 0
+    except:
+        return True
+
+
 class COCODataset(torchvision.datasets.coco.CocoDetection):
     def __init__(
         self, ann_file, root, remove_images_without_annotations, transforms=None
@@ -29,7 +36,7 @@ class COCODataset(torchvision.datasets.coco.CocoDetection):
                 if all(
                     any(o <= 1 for o in obj["bbox"][2:])
                     for obj in anno
-                    if obj["iscrowd"] == 0
+                    if iscrowd_is_0(obj)
                 ):
                     ids_to_remove.append(img_id)

@@ -51,7 +58,7 @@ class COCODataset(torchvision.datasets.coco.CocoDetection):

         # filter crowd annotations
         # TODO might be better to add an extra field
-        anno = [obj for obj in anno if obj["iscrowd"] == 0]
+        anno = [obj for obj in anno if iscrowd_is_0(obj)]

         boxes = [obj["bbox"] for obj in anno]
         boxes = torch.as_tensor(boxes).reshape(-1, 4)  # guard against no boxes

diff --git a/c2py/maskrcnn_benchmark/engine/trainer.py b/c2py/maskrcnn_benchmark/engine/trainer.py
index 38a9e52..85e281a 100644
--- a/c2py/maskrcnn_benchmark/engine/trainer.py
+++ b/c2py/maskrcnn_benchmark/engine/trainer.py
@@ -8,6 +8,10 @@ import torch.distributed as dist

 from maskrcnn_benchmark.utils.comm import get_world_size
 from maskrcnn_benchmark.utils.metric_logger import MetricLogger
+from maskrcnn_benchmark.utils.tensorboard_logger import TensorboardLogger
+
+ENABLE_TENSORBOARD = True
+


 def reduce_loss_dict(loss_dict):
@@ -47,7 +51,10 @@ def do_train(
 ):
     logger = logging.getLogger("maskrcnn_benchmark.trainer")
     logger.info("Start training")
-    meters = MetricLogger(delimiter="  ")
+    if ENABLE_TENSORBOARD:
+        meters = TensorboardLogger()
+    else:
+        meters = MetricLogger(delimiter="  ")
     max_iter = len(data_loader)
     start_iter = arguments["iteration"]
     model.train()
@@ -78,7 +85,14 @@ def do_train(

         batch_time = time.time() - end
         end = time.time()
-        meters.update(time=batch_time, data=data_time)
+        if ENABLE_TENSORBOARD:
+            meters.update(
+                loss=losses_reduced,
+                time=batch_time,
+                data=data_time,
+                **loss_dict_reduced)
+        else:
+            meters.update(time=batch_time, data=data_time)

         eta_seconds = meters.time.global_avg * (max_iter - iteration)
         eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
diff --git a/c2py/maskrcnn_benchmark/utils/tensorboard_logger.py b/c2py/maskrcnn_benchmark/utils/tensorboard_logger.py
new file mode 100644
index 0000000..0ce637a
--- /dev/null
+++ b/c2py/maskrcnn_benchmark/utils/tensorboard_logger.py
@@ -0,0 +1,22 @@
+import torch
+from maskrcnn_benchmark.utils.metric_logger import MetricLogger
+
+from tensorboardX import SummaryWriter
+
+
+class TensorboardLogger(MetricLogger):
+    def __init__(self, log_dir='log', start_iter=0, delimiter="\t"):
+        super(TensorboardLogger, self).__init__(delimiter)
+        self.iteration = start_iter
+        self.writer = SummaryWriter(log_dir)
+
+    def update(self, **kwargs):
+        super(TensorboardLogger, self).update(**kwargs)
+        if self.writer:
+            for k, v in kwargs.items():
+                if isinstance(v, torch.Tensor):
+                    v = v.item()
+                assert isinstance(v, (float, int))
+                self.writer.add_scalar(k, v, self.iteration)
+        self.iteration += 1