YOLO webcam example - add assert for webcam load (#202) (#204)

* YOLO webcam example - add assert for webcam load * update readme to note other options for annotate * formatting
neuralmagic · Apr 29, 2021 · f7dbd57 · f7dbd57
1 parent c47bfb8
commit f7dbd57
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/integrations/ultralytics/README.md b/integrations/ultralytics/README.md
@@ -151,9 +151,14 @@ python annotate.py \
     zoo:cv/detection/yolo_v3-spp/pytorch/ultralytics/coco/pruned_quant-aggressive_94 \
     --source 0 \
     --quantized-inputs \
-    --image-shape 416 416
+    --image-shape 416 416 \
+    --no-save  # webcam only
 ```
 
+In addition to webcam `--source` can take a path to a `.jpg` file, directory or glog pat
+of `.jpg` files, or path to a `.mp4` video file.  If source is an integer and no
+corresponding webcam is available, an exception will be raised.
+
 
 ### Benchmarking
 `benchmarking.py` is a script for benchmarking sparsified and quantized YOLOv3

diff --git a/integrations/ultralytics/deepsparse/deepsparse_utils.py b/integrations/ultralytics/deepsparse/deepsparse_utils.py
@@ -199,7 +199,10 @@ def __iter__(self) -> Iterator[Tuple[numpy.ndarray, numpy.ndarray]]:
                 self._stream.release()
                 cv2.destroyAllWindows()
                 break
-            _, frame = self._stream.read()
+            loaded, frame = self._stream.read()
+
+            assert loaded, f"Could not load image from webcam {self._camera}"
+
             frame = cv2.flip(frame, 1)  # flip left-right
             yield load_image(frame, image_size=self._image_size)