From 9f4d11379bb931586c1f51c1d85c6fac9fc37eda Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 20 Mar 2022 14:18:05 +0100 Subject: [PATCH] Fix incomplete URL substring sanitation (#7056) Resolves code scanning alert in https://github.com/ultralytics/yolov5/issues/7055 --- utils/datasets.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/datasets.py b/utils/datasets.py index 5ce6d607fb7a..8627344af7b4 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -15,6 +15,7 @@ from multiprocessing.pool import Pool, ThreadPool from pathlib import Path from threading import Thread +from urllib.parse import urlparse from zipfile import ZipFile import cv2 @@ -301,7 +302,7 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True): for i, s in enumerate(sources): # index, source # Start thread to read frames from video stream st = f'{i + 1}/{n}: {s}... ' - if 'youtube.com/' in s or 'youtu.be/' in s: # if source is YouTube video + if urlparse(s).hostname in ('youtube.com', 'youtu.be'): # if source is YouTube video check_requirements(('pafy', 'youtube_dl==2020.12.2')) import pafy s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL