Improve code quality: proper logging, error handling, and optional de…

datasets/prepare_panoptic_fpn.py

@@ -4,6 +4,7 @@
 
 import functools
 import json
+import logging
 import multiprocessing as mp
 import numpy as np
 import os
@@ -67,14 +68,17 @@ def iter_annotations():
             output = os.path.join(sem_seg_root, file_name)
             yield input, output, segments
 
-    print("Start writing to {} ...".format(sem_seg_root))
+    if not logging.getLogger().hasHandlers():
+        logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+    logger.info("Start writing to %s ...", sem_seg_root)
     start = time.time()
     pool.starmap(
         functools.partial(_process_panoptic_to_semantic, id_map=id_map),
         iter_annotations(),
         chunksize=100,
     )
-    print("Finished. time: {:.2f}s".format(time.time() - start))
+    logger.info("Finished. time: %.2fs", time.time() - start)
 
 
 if __name__ == "__main__":

demo/demo.py

@@ -104,7 +104,8 @@ def main() -> None:
     if args.input:
         if len(args.input) == 1:
             args.input = glob.glob(os.path.expanduser(args.input[0]))
-            assert args.input, "The input path(s) was not found"
+            if not args.input:
+                raise FileNotFoundError("The input path(s) was not found")
         for path in tqdm.tqdm(args.input, disable=not args.output):
             # use PIL, to be consistent with evaluation
             img = read_image(path, format="BGR")
@@ -124,10 +125,10 @@ def main() -> None:
 
             if args.output:
                 if os.path.isdir(args.output):
-                    assert os.path.isdir(args.output), args.output
                     out_filename = os.path.join(args.output, os.path.basename(path))
                 else:
-                    assert len(args.input) == 1, "Please specify a directory with args.output"
+                    if len(args.input) != 1:
+                        raise ValueError("Please specify a directory with args.output")
                     out_filename = args.output
                 visualized_output.save(out_filename)
             else:
@@ -136,8 +137,10 @@ def main() -> None:
                 if cv2.waitKey(0) == 27:
                     break  # esc to quit
     elif args.webcam:
-        assert args.input is None, "Cannot have both --input and --webcam!"
-        assert args.output is None, "output not yet supported with --webcam!"
+        if args.input is not None:
+            raise ValueError("Cannot have both --input and --webcam!")
+        if args.output is not None:
+            raise ValueError("output not yet supported with --webcam!")
         cam = cv2.VideoCapture(0)
         for vis in tqdm.tqdm(demo.run_on_video(cam)):
             cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
@@ -164,7 +167,8 @@ def main() -> None:
                 output_fname = os.path.splitext(output_fname)[0] + file_ext
             else:
                 output_fname = args.output
-            assert not os.path.isfile(output_fname), output_fname
+            if os.path.isfile(output_fname):
+                raise ValueError(f"Output file already exists: {output_fname}")
             output_file = cv2.VideoWriter(
                 filename=output_fname,
                 # some installation of opencv may not support x264 (due to its license),
@@ -174,7 +178,8 @@ def main() -> None:
                 frameSize=(width, height),
                 isColor=True,
             )
-        assert os.path.isfile(args.video_input)
+        if not os.path.isfile(args.video_input):
+            raise FileNotFoundError(f"Video input file not found: {args.video_input}")
         for vis_frame in tqdm.tqdm(demo.run_on_video(video), total=num_frames):
             if args.output:
                 output_file.write(vis_frame)

detectron2/utils/visualizer.py

@@ -5,12 +5,18 @@
 import numpy as np
 from enum import Enum, unique
 import cv2
-import matplotlib as mpl
-import matplotlib.colors as mplc
-import matplotlib.figure as mplfigure
+try:
+    import matplotlib as mpl
+    import matplotlib.colors as mplc
+    import matplotlib.figure as mplfigure
+    from matplotlib.backends.backend_agg import FigureCanvasAgg
+except ImportError:
+    mpl = None
+    mplc = None
+    mplfigure = None
+    FigureCanvasAgg = None
 import pycocotools.mask as mask_util
 import torch
-from matplotlib.backends.backend_agg import FigureCanvasAgg
 from PIL import Image
 
 from detectron2.data import MetadataCatalog

setup.py

@@ -164,7 +164,6 @@ def get_model_zoo_configs() -> List[str]:
         # In general, avoid adding dependencies that are not pure-python because they are not
         # guaranteed to be installable by `pip install` on all platforms.
         "Pillow>=7.1",  # or use pillow-simd for better performance
-        "matplotlib",  # TODO move it to optional after we add opencv visualization
         "pycocotools>=2.0.2",  # corresponds to https://github.com/ppwwyyxx/cocoapi
         # Do not add opencv here. Just like pytorch, user should install
         # opencv themselves, preferrably by OS's package manager, or by
@@ -205,6 +204,8 @@ def get_model_zoo_configs() -> List[str]:
             "psutil",
             "panopticapi @ https://github.com/cocodataset/panopticapi/archive/master.zip",
         ],
+        # visualization dependencies
+        "viz": ["matplotlib"],
         # dev dependencies. Install them by `pip install 'detectron2[dev]'`
         "dev": [
             "flake8==3.8.1",

tools/train_net.py

@@ -123,6 +123,9 @@ def setup(args):
 
 def main(args):
     cfg = setup(args)
+
+    logger = logging.getLogger(__name__)
+    logger.info("Command Line Args: %s", args)
 
     if args.eval_only:
         model = Trainer.build_model(cfg)
@@ -152,7 +155,6 @@ def main(args):
 
 def invoke_main() -> None:
     args = default_argument_parser().parse_args()
-    print("Command Line Args:", args)
     launch(
         main,
         args.num_gpus,

tools/visualize_data.py

@@ -57,12 +57,12 @@ def main() -> None:
 
     def output(vis, fname):
         if args.show:
-            print(fname)
+            logger.info(fname)
             cv2.imshow("window", vis.get_image()[:, :, ::-1])
             cv2.waitKey()
         else:
             filepath = os.path.join(dirname, fname)
-            print("Saving to {} ...".format(filepath))
+            logger.info("Saving to %s ...", filepath)
             vis.save(filepath)
 
     scale = 1.0