From da49d53946bea637ffb72bb295ed1bd9fe738cc0 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 16 Feb 2021 13:56:47 -0800
Subject: [PATCH] Add check_imshow() (#2231)

* Add check_imshow()

* Update general.py

* Update general.py
---
 detect.py        |  8 ++++----
 utils/general.py | 13 +++++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/detect.py b/detect.py
index 3f1d6c521b67..22bf21b4c825 100644
--- a/detect.py
+++ b/detect.py
@@ -9,8 +9,8 @@
 
 from models.experimental import attempt_load
 from utils.datasets import LoadStreams, LoadImages
-from utils.general import check_img_size, check_requirements, non_max_suppression, apply_classifier, scale_coords, \
-    xyxy2xywh, strip_optimizer, set_logging, increment_path
+from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
+    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
 from utils.plots import plot_one_box
 from utils.torch_utils import select_device, load_classifier, time_synchronized
 
@@ -45,7 +45,7 @@ def detect(save_img=False):
     # Set Dataloader
     vid_path, vid_writer = None, None
     if webcam:
-        view_img = True
+        view_img = check_imshow()
         cudnn.benchmark = True  # set True to speed up constant image size inference
         dataset = LoadStreams(source, img_size=imgsz, stride=stride)
     else:
@@ -118,7 +118,7 @@ def detect(save_img=False):
             # Stream results
             if view_img:
                 cv2.imshow(str(p), im0)
-                cv2.waitKey(1) # 1 millisecond
+                cv2.waitKey(1)  # 1 millisecond
 
             # Save results (image with detections)
             if save_img:
diff --git a/utils/general.py b/utils/general.py
index 24807483f5f4..2d3e83ede35e 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -95,6 +95,19 @@ def check_img_size(img_size, s=32):
     return new_size
 
 
+def check_imshow():
+    # Check if environment supports image displays
+    try:
+        cv2.imshow('test', np.zeros((1, 1, 3)))
+        cv2.waitKey(1)
+        cv2.destroyAllWindows()
+        cv2.waitKey(1)
+        return True
+    except Exception as e:
+        print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image previews\n{e}')
+        return False
+
+
 def check_file(file):
     # Search for file if not found
     if os.path.isfile(file) or file == '':