Replaced print statements with logging library (#95)

test/common_utils.py

@@ -17,6 +17,8 @@
 import numpy as np
 from PIL import Image
 
+import logging
+logger = logging.getLogger(__name__)
 
 def set_rng_seed(seed):
     torch.manual_seed(seed)
@@ -130,7 +132,7 @@ def assertExpected(self, output, subname=None, prec=None, strip_suffix=None):
 
         if ACCEPT:
             filename = {os.path.basename(expected_file)}
-            print("Accepting updated output for {}:\n\n{}".format(filename, output))
+            logger.info("Accepting updated output for {}:\n\n{}".format(filename, output))
             torch.save(output, expected_file)
             MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
             binary_size = os.path.getsize(expected_file)

yolort/data/_helper.py

@@ -9,6 +9,8 @@
 from .coco import COCODetection
 from .transforms import collate_fn, default_train_transforms, default_val_transforms
 
+import logging
+logger = logging.getLogger(__name__)
 
 def get_coco_api_from_dataset(dataset):
     for _ in range(10):
@@ -34,18 +36,18 @@ def prepare_coco128(
         dirname (str): the directory name of coco128 dataset. Default: 'coco128'.
     """
     if not data_path.is_dir():
-        print(f'Create a new directory: {data_path}')
+        logger.info(f'Create a new directory: {data_path}')
         data_path.mkdir(parents=True, exist_ok=True)
 
     zip_path = data_path / 'coco128.zip'
     coco128_url = 'https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip'
     if not zip_path.is_file():
-        print(f'Downloading coco128 datasets form {coco128_url}')
+        logger.info(f'Downloading coco128 datasets form {coco128_url}')
         torch.hub.download_url_to_file(coco128_url, zip_path, hash_prefix='a67d2887')
 
     coco128_path = data_path / dirname
     if not coco128_path.is_dir():
-        print(f'Unzipping dataset to {coco128_path}')
+        logger.info(f'Unzipping dataset to {coco128_path}')
         with ZipFile(zip_path, 'r') as zip_obj:
             zip_obj.extractall(data_path)
 

yolort/models/experimental.py

@@ -7,6 +7,9 @@
 
 from .common import Conv, DWConv
 
+import logging
+logger = logging.getLogger(__name__)
+
 
 class CrossConv(nn.Module):
     # Cross Convolution Downsample
@@ -126,7 +129,7 @@ def attempt_load(weights, map_location=None):
     if len(model) == 1:
         return model[-1]  # return model
     else:
-        print('Ensemble created with %s\n' % weights)
+        logger.info('Ensemble created with %s\n' % weights)
         for k in ['names', 'stride']:
             setattr(model, k, getattr(model[-1], k))
         return model  # return ensemble

yolort/utils/image_utils.py

@@ -16,6 +16,9 @@
 
 from typing import Optional
 
+import logging
+logger = logging.getLogger(__name__)
+
 
 def plot_one_box(box, img, color=None, label=None, line_thickness=None):
     # Plots one bounding box on image img
@@ -225,7 +228,7 @@ def overlay_boxes(detections, path, time_consume, args):
                     )
 
         # Print inference time
-        print('%sDone. (%.3fs)' % (det_logs, time_consume))
+        logger.info('%sDone. (%.3fs)' % (det_logs, time_consume))
 
         # Save results (image with detections)
         if args.save_img: