Fix incomplete URL substring sanitation (#7056)
Browse filesResolves code scanning alert in https://github.com/ultralytics/yolov5/issues/7055
- utils/datasets.py +2 -1
utils/datasets.py
CHANGED
@@ -15,6 +15,7 @@ from itertools import repeat
|
|
15 |
from multiprocessing.pool import Pool, ThreadPool
|
16 |
from pathlib import Path
|
17 |
from threading import Thread
|
|
|
18 |
from zipfile import ZipFile
|
19 |
|
20 |
import cv2
|
@@ -301,7 +302,7 @@ class LoadStreams:
|
|
301 |
for i, s in enumerate(sources): # index, source
|
302 |
# Start thread to read frames from video stream
|
303 |
st = f'{i + 1}/{n}: {s}... '
|
304 |
-
if 'youtube.com
|
305 |
check_requirements(('pafy', 'youtube_dl==2020.12.2'))
|
306 |
import pafy
|
307 |
s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL
|
|
|
15 |
from multiprocessing.pool import Pool, ThreadPool
|
16 |
from pathlib import Path
|
17 |
from threading import Thread
|
18 |
+
from urllib.parse import urlparse
|
19 |
from zipfile import ZipFile
|
20 |
|
21 |
import cv2
|
|
|
302 |
for i, s in enumerate(sources): # index, source
|
303 |
# Start thread to read frames from video stream
|
304 |
st = f'{i + 1}/{n}: {s}... '
|
305 |
+
if urlparse(s).hostname in ('youtube.com', 'youtu.be'): # if source is YouTube video
|
306 |
check_requirements(('pafy', 'youtube_dl==2020.12.2'))
|
307 |
import pafy
|
308 |
s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL
|