Set OpenCV as optional (#286)

yolort/utils/image_utils.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 from typing import Optional
 
-import cv2
 import matplotlib.pyplot as plt
 import numpy as np
 import requests
@@ -13,6 +12,11 @@
 from torch import Tensor
 from torchvision.ops.boxes import box_convert
 
+try:
+    import cv2
+except ImportError:
+    cv2 = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -80,10 +84,7 @@ def hex2rgb(h):
     return [hex2rgb(h) for h in plt.rcParams["axes.prop_cycle"].by_key()["color"]]
 
 
-def get_image_from_url(
-    url: str,
-    flags: int = cv2.IMREAD_COLOR,
-) -> np.ndarray:
+def get_image_from_url(url: str, flags: int = 1) -> np.ndarray:
     """
     Generates an image directly from an URL
 

yolort/v5/utils/augmentations.py

@@ -7,9 +7,13 @@
 import math
 import random
 
-import cv2
 import numpy as np
 
+try:
+    import cv2
+except ImportError:
+    cv2 = None
+
 from .general import colorstr, segment2box, resample_segments, check_version
 from .metrics import bbox_ioa
 

yolort/v5/utils/general.py

@@ -16,14 +16,18 @@
 import urllib
 from pathlib import Path
 
-import cv2
 import numpy as np
 import pandas as pd
 import pkg_resources as pkg
 import torch
 import torchvision
 import yaml
 
+try:
+    import cv2
+except ImportError:
+    cv2 = None
+
 from .metrics import box_iou, fitness
 
 # Settings
@@ -32,7 +36,8 @@
 np.set_printoptions(linewidth=320, formatter={"float_kind": "{:11.5g}".format})
 pd.options.display.max_columns = 10
 # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
-cv2.setNumThreads(0)
+if cv2 is not None:
+    cv2.setNumThreads(0)
 os.environ["NUMEXPR_MAX_THREADS"] = str(min(os.cpu_count(), 8))  # NumExpr max threads
 
 FILE = Path(__file__).resolve()

yolort/v5/utils/plots.py

@@ -8,14 +8,18 @@
 from copy import copy
 from pathlib import Path
 
-import cv2
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import torch
 from PIL import Image, ImageDraw, ImageFont
 
+try:
+    import cv2
+except ImportError:
+    cv2 = None
+
 from .general import (
     LOGGER,
     Timeout,