glenn-jocher commited on
Commit
9f4d113
·
unverified ·
1 Parent(s): 6843ea5

Fix incomplete URL substring sanitation (#7056)

Browse files

Resolves code scanning alert in https://github.com/ultralytics/yolov5/issues/7055

Files changed (1) hide show
  1. 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/' in s or 'youtu.be/' in s: # if source is YouTube video
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