diff --git a/supervision/assets/downloader.py b/supervision/assets/downloader.py index b5312d9c..2d494dca 100644 --- a/supervision/assets/downloader.py +++ b/supervision/assets/downloader.py @@ -1,3 +1,4 @@ +import os from functools import partial from hashlib import md5 from os import remove as os_remove @@ -22,31 +23,33 @@ except ImportError: ) -# md5 check function -def _check_md5(filename: str, orig_md5: str) -> bool: +def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool: """ - filename: str, A string representing the path to the file. - orig_md5: str, A string representing the original md5 hash. + Check if the MD5 hash of a file matches the original hash. + + Parameters: + filename (str): The path to the file to be checked as a string. + original_md5_hash (str): The original MD5 hash to compare against. + + Returns: + bool: True if the hashes match, False otherwise. """ - if not path_exists(filename): - return False - with open(filename, "rb") as file_to_check: - # read contents of the file - data = file_to_check.read() - # pipe contents of the file through - md5_returned = md5(data).hexdigest() - # Return True if the computed hash matches the original one - if md5_returned == orig_md5: - return True + if not os.path.exists(filename): return False + with open(filename, "rb") as file: + file_contents = file.read() + computed_md5_hash = md5(file_contents).hexdigest() + + return computed_md5_hash == original_md5_hash + def download_assets(asset_name: Union[VideoAssets, str]) -> str: """ asset_name: VIDEO_ASSETS, name of the file to download """ - filename = f"{asset_name}" + filename = asset_name.value if isinstance(asset_name, VideoAssets) else asset_name if not path_exists(filename) and asset_name in VIDEO_ASSETS: print(f"Downloading {filename} assets \n") res: Response = get( @@ -66,19 +69,19 @@ def download_assets(asset_name: Union[VideoAssets, str]) -> str: desc = "(Unknown total file size)" if file_size == 0 else "" res.raw.read = partial( res.raw.read, decode_content=True - ) # Decompress if needed + ) with tqdm.wrapattr( res.raw, "read", total=file_size, desc=desc, - colour="a351fb", + colour="#a351fb", ) as r_raw: with folder_path.open("wb") as f: copyfileobj(r_raw, f) elif path_exists(filename): - if not _check_md5(filename, VIDEO_ASSETS[asset_name][2]): + if not is_md5_hash_matching(filename, VIDEO_ASSETS[asset_name][2]): print("File corrupted. Re-downloading... \n") os_remove(filename) download_assets(asset_name) @@ -86,14 +89,7 @@ def download_assets(asset_name: Union[VideoAssets, str]) -> str: print(f"{filename} asset download complete. \n") else: raise ValueError( - "Invalid asset type. It should be one of the following: \n" - "vehicles.mp4 \n" - "vehicles-2.mp4 \n" - "milk-video-1.mp4 \n" - "grocery-store.mp4 \n" - "subway.mp4 \n" - "market-square.mp4 \n" - "people-walking-bw.mp4 \n" + f"Invalid asset. It should be one of the following: {VideoAssets.list()}." ) return filename