diff --git a/app.py b/app.py index 0f250cf..5c48fe1 100644 --- a/app.py +++ b/app.py @@ -68,8 +68,6 @@ def input(): form = request.form input_vars = [ - 'local', - 'playlist', 'graphs', 'minimap', 'lightweight' @@ -80,7 +78,7 @@ def input(): for var in input_vars: input_args[var] = var in form - run_video_thread = threading.Thread( + threading.Thread( target=lava.execute, name="tracker", kwargs=input_args diff --git a/assets/lava.py b/assets/lava.py index 435b962..61d57e1 100644 --- a/assets/lava.py +++ b/assets/lava.py @@ -6,11 +6,12 @@ from numpy.linalg import norm from sympy.geometry import Point, Circle, intersection, Line from sympy import N -from youtubesearchpython import Playlist, Video, StreamURLFetcher +from assets.utils import utils +from assets.utils.ytHelper import is_valid_youtube_url, parse_youtube_url +from assets.utils.videoFileHelper import parse_local_files from assets.graphing import GraphsOperator -from assets import utils, constants - +import assets.utils.constants as constants class LAVA(GraphsOperator): def __init__(self): @@ -29,37 +30,28 @@ def __init__(self): def execute(self, url="", - local=False, - playlist=False, output_file='positions.csv', - skip=0, minimap=False, graphs=False, lightweight=True): df = self.gather_data( url, - local, output_file, - playlist, - skip, minimap, lightweight ) if graphs is True: self.draw_graphs(df, df.video.iloc[0]) - + print(f'Video {df.video.iloc[0]} complete') def gather_data(self, url="", - local=False, output_file='positions.csv', - playlist=False, - skip=0, minimap=False, lightweight=True): - videos = self.parse_url(url, local, playlist, skip) + videos = self.parse_url(url) full_data = pd.DataFrame() self.lightweight = lightweight @@ -119,9 +111,8 @@ def gather_data(self, cv2.destroyAllWindows() return full_data - - def gather_info(self, url="", local=False, playlist=False, skip=0): - videos = self.parse_url(url, local, playlist, skip) + def gather_info(self, url=""): + videos = self.parse_url(url) full_data = {} @@ -150,49 +141,11 @@ def gather_info(self, url="", local=False, playlist=False, skip=0): full_data[video] = champs return full_data - def parse_url(self, url="", local=False, playlist=False, skip=0): - videos = {} - - if local is False: - fetcher = StreamURLFetcher() - - if playlist is True: - playlist_data = Playlist.getVideos(url) - videos_list = [x['link'] for x in playlist_data['videos']] - videos_list = videos_list[skip:] - else: - if isinstance(url, list): - videos_list = url - else: - videos_list = [url] - - for video_url in videos_list: - video = Video.get(video_url) - url = fetcher.get(video, 22) - - video = utils.clean_for_directory(video['title']) - - videos[video] = url - # If pulling from local + def parse_url(self, url=""): + if is_valid_youtube_url(url): + return parse_youtube_url(url) else: - if playlist is True: - videos_list = os.listdir('input') - videos_list.remove('.gitkeep') - videos_list = videos_list[skip:] - else: - if isinstance(url, list): - videos_list = url - else: - videos_list = [url] - - for video_file in videos_list: - video_path = video_file - - video = os.path.splitext(video_file)[0] - - videos[video] = video_path - - return videos + return parse_local_files(url) def headers(self): headers_path = os.path.join( diff --git a/assets/utils/__init__.py b/assets/utils/__init__.py new file mode 100644 index 0000000..2f5c78e --- /dev/null +++ b/assets/utils/__init__.py @@ -0,0 +1,2 @@ +import pkgutil +__path__ = pkgutil.extend_path(__path__, __name__) \ No newline at end of file diff --git a/assets/constants.py b/assets/utils/constants.py similarity index 99% rename from assets/constants.py rename to assets/utils/constants.py index cc9dca0..8a678dd 100644 --- a/assets/constants.py +++ b/assets/utils/constants.py @@ -2,6 +2,7 @@ import cv2 + tracking_folder = os.path.join( "assets", "tracking") diff --git a/assets/utils.py b/assets/utils/utils.py similarity index 87% rename from assets/utils.py rename to assets/utils/utils.py index 6758ff0..4a53fec 100644 --- a/assets/utils.py +++ b/assets/utils/utils.py @@ -4,15 +4,6 @@ import cv2 -def clean_for_directory(video_title): - for ch in ['*', '.', '"', '/', '\\', ':', ';', '|', ',']: - if ch in video_title: - video_title = video_title.replace(ch, '') - video_title = video_title.replace(' ', '_') - - return video_title - - def get_header_borders(frame_shape): frame_height, frame_width, _ = frame_shape header_height = frame_height // 15 diff --git a/assets/utils/video_file_helper.py b/assets/utils/video_file_helper.py new file mode 100644 index 0000000..78d0aa3 --- /dev/null +++ b/assets/utils/video_file_helper.py @@ -0,0 +1,31 @@ +import os + +VALID_VIDEO_FORMATS = ["WEBM", "MPG", "MP2", "MPEG", "MPE", "MPV", "OGG", "MP4", "M4P", "M4V", "AVI", "WMV", "MOV", "QT", "FLV", "SWF", "AVCHD"] + +def clean_for_directory(video_title : str) -> str: + for ch in ['*', '.', '"', '/', '\\', ':', ';', '|', ',']: + if ch in video_title: + video_title = video_title.replace(ch, '') + video_title = video_title.replace(' ', '_') + + return video_title + +def parse_local_files(url : str) -> dict or None: + videos = {} + print(url) + if os.path.isdir(url): + for file in os.listdir(url): + if file.split('.')[-1].upper() in VALID_VIDEO_FORMATS: + # get the file's name + video_name = clean_for_directory(file) + videos[video_name] = os.path.join(url, file) + else: + print(f'Invalid file format: {file}') + elif os.path.isfile(url): + if url.split('.')[-1].upper() in VALID_VIDEO_FORMATS: + video_name = clean_for_directory(url) + videos[video_name] = url + return videos + else: + print(f'Invalid file format: {url}') + return None \ No newline at end of file diff --git a/assets/utils/youtube_helper.py b/assets/utils/youtube_helper.py new file mode 100644 index 0000000..b495ed9 --- /dev/null +++ b/assets/utils/youtube_helper.py @@ -0,0 +1,27 @@ +import yt_dlp +from assets.utils.video_file_helper import clean_for_directory + +def is_valid_youtube_url(url) -> bool: + return 'youtube' in url and '/watch?v=' in url + +def is_youtube_playlist(url) -> bool: + return '&list=' in url or '/playlist?list=' in url + +def parse_youtube_url(url : str) -> dict: + videos = {} + ydl_opts = { + 'format': '22', + 'ignoreerrors': True, + 'no_warnings': True + } + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + info = ydl.extract_info(url, download=False) + video_info = ydl.sanitize_info(info) + if video_info['_type'] == 'playlist': + for video in video_info['entries']: + video_name = clean_for_directory(video['title']) + videos[video_name] = video_info['formats'][-1]['url'] + else: + video_name = clean_for_directory(video_info['title']) + videos[video_name] = video_info['formats'][-1]['url'] + return videos diff --git a/assets/webui_templates/input.html b/assets/webui_templates/input.html index c0a62e2..f9a6a41 100644 --- a/assets/webui_templates/input.html +++ b/assets/webui_templates/input.html @@ -49,27 +49,12 @@