fix issue on yolov3 tutorial

DangerStep · Apr 26, 2022 · af7dff5 · af7dff5
1 parent 3105f40
commit af7dff5
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/machine-learning/object-detection/live_yolo_opencv.py b/machine-learning/object-detection/live_yolo_opencv.py
@@ -16,7 +16,11 @@
 net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
 
 ln = net.getLayerNames()
-ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+try:
+    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+except IndexError:
+    # in case getUnconnectedOutLayers() returns 1D array when CUDA isn't available
+    ln = [ln[i - 1] for i in net.getUnconnectedOutLayers()]
 
 cap = cv2.VideoCapture(0)
 

diff --git a/machine-learning/object-detection/read_video.py b/machine-learning/object-detection/read_video.py
@@ -17,7 +17,11 @@
 net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
 
 ln = net.getLayerNames()
-ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+try:
+    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+except IndexError:
+    # in case getUnconnectedOutLayers() returns 1D array when CUDA isn't available
+    ln = [ln[i - 1] for i in net.getUnconnectedOutLayers()]
 # read the file from the command line
 video_file = sys.argv[1]
 cap = cv2.VideoCapture(video_file)

diff --git a/machine-learning/object-detection/yolo_opencv.py b/machine-learning/object-detection/yolo_opencv.py
@@ -37,7 +37,11 @@
 
 # get all the layer names
 ln = net.getLayerNames()
-ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+try:
+    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+except IndexError:
+    # in case getUnconnectedOutLayers() returns 1D array when CUDA isn't available
+    ln = [ln[i - 1] for i in net.getUnconnectedOutLayers()]
 # feed forward (inference) and get the network output
 # measure how much it took in seconds
 start = time.perf_counter()