diff --git a/docs/assets.md b/docs/assets.md index 7135f33f..7780b903 100644 --- a/docs/assets.md +++ b/docs/assets.md @@ -4,7 +4,7 @@ comments: true # Assets -Supervision offers an assets download utility that allows you to download video files +Supervision offers an assets download utility that allows you to download image and video files that you can use in your demos.
@@ -18,3 +18,9 @@ that you can use in your demos.
:::supervision.assets.list.VideoAssets + +
+

ImageAssets

+
+ +:::supervision.assets.list.ImageAssets diff --git a/src/supervision/assets/__init__.py b/src/supervision/assets/__init__.py index 4c48bab6..3ed213f5 100644 --- a/src/supervision/assets/__init__.py +++ b/src/supervision/assets/__init__.py @@ -1,4 +1,4 @@ from supervision.assets.downloader import download_assets -from supervision.assets.list import VideoAssets +from supervision.assets.list import ImageAssets, VideoAssets -__all__ = ["VideoAssets", "download_assets"] +__all__ = ["ImageAssets", "VideoAssets", "download_assets"] diff --git a/src/supervision/assets/downloader.py b/src/supervision/assets/downloader.py index e535a95a..f537dd18 100644 --- a/src/supervision/assets/downloader.py +++ b/src/supervision/assets/downloader.py @@ -8,7 +8,7 @@ from shutil import copyfileobj from requests import get from tqdm.auto import tqdm -from supervision.assets.list import VIDEO_ASSETS, VideoAssets +from supervision.assets.list import MEDIA_ASSETS, Assets def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool: @@ -35,7 +35,7 @@ def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool: return computed_md5_hash.hexdigest() == original_md5_hash -def download_assets(asset_name: VideoAssets | str) -> str: +def download_assets(asset_name: Assets | str) -> str: """ Download a specified asset if it doesn't already exist or is corrupted. @@ -47,42 +47,44 @@ def download_assets(asset_name: VideoAssets | str) -> str: Example: ```pycon - >>> from supervision.assets import download_assets, VideoAssets + >>> from supervision.assets import download_assets, ImageAssets, VideoAssets >>> download_assets(VideoAssets.VEHICLES) # doctest: +SKIP 'vehicles.mp4' + >>> download_assets(ImageAssets.PEOPLE_WALKING) # doctest: +SKIP + 'people-walking.jpg' + ``` """ - filename = asset_name.value if isinstance(asset_name, VideoAssets) else asset_name + filename = asset_name.filename if isinstance(asset_name, Assets) else asset_name - if not Path(filename).exists() and filename in VIDEO_ASSETS: - print(f"Downloading {filename} assets \n") - response = get( - VIDEO_ASSETS[filename][0], stream=True, allow_redirects=True, timeout=30 - ) - response.raise_for_status() + if filename in MEDIA_ASSETS: + if not Path(filename).exists(): + print(f"Downloading {filename} assets \n") + response = get( + MEDIA_ASSETS[filename][0], stream=True, allow_redirects=True, timeout=30 + ) + response.raise_for_status() - file_size = int(response.headers.get("Content-Length", 0)) - folder_path = Path(filename).expanduser().resolve() - folder_path.parent.mkdir(parents=True, exist_ok=True) + file_size = int(response.headers.get("Content-Length", 0)) + folder_path = Path(filename).expanduser().resolve() + folder_path.parent.mkdir(parents=True, exist_ok=True) - with tqdm.wrapattr( - response.raw, "read", total=file_size, desc="", colour="#a351fb" - ) as raw_resp: - with folder_path.open("wb") as file: - copyfileobj(raw_resp, file) - - elif Path(filename).exists(): - if not is_md5_hash_matching(filename, VIDEO_ASSETS[filename][1]): - print("File corrupted. Re-downloading... \n") - os.remove(filename) - return download_assets(filename) - - print(f"{filename} asset download complete. \n") + with tqdm.wrapattr( + response.raw, "read", total=file_size, desc="", colour="#a351fb" + ) as raw_resp: + with folder_path.open("wb") as file: + copyfileobj(raw_resp, file) + else: + if not is_md5_hash_matching(filename, MEDIA_ASSETS[filename][1]): + print("File corrupted. Re-downloading... \n") + os.remove(filename) + return download_assets(filename) + print(f"{filename} asset download complete. \n") else: - valid_assets = ", ".join(asset.value for asset in VideoAssets) + valid_assets = ", ".join(filename for filename in MEDIA_ASSETS.keys()) raise ValueError( f"Invalid asset. It should be one of the following: {valid_assets}." ) diff --git a/src/supervision/assets/list.py b/src/supervision/assets/list.py index 0d6d0ae0..c760895d 100644 --- a/src/supervision/assets/list.py +++ b/src/supervision/assets/list.py @@ -1,14 +1,31 @@ from enum import Enum BASE_VIDEO_URL = "https://media.roboflow.com/supervision/video-examples/" +BASE_IMAGE_URL = "https://media.roboflow.com/supervision/image-examples/" -class VideoAssets(Enum): +class Assets(Enum): + filename: str + md5_hash: str + + def __new__(cls, filename: str, md5_hash: str) -> "Assets": + obj = object.__new__(cls) + obj._value_ = filename + obj.filename = filename + obj.md5_hash = md5_hash + return obj + + @classmethod + def list(cls) -> list[str]: + return [asset.filename for asset in cls] + + +class VideoAssets(Assets): """ - Each member of this enum represents a video asset. The value associated with each - member is the filename of the video. + Each member of this class represents a video asset. The value associated with each + member has a filename and hash of the video. File names and links can be seen below. - | Enum Member | Video Filename | Video URL | + | Asset | Video Filename | Video URL | |------------------------|----------------------------|---------------------------------------------------------------------------------------| | `VEHICLES` | `vehicles.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/vehicles.mp4) | | `MILK_BOTTLING_PLANT` | `milk-bottling-plant.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/milk-bottling-plant.mp4) | @@ -22,61 +39,44 @@ class VideoAssets(Enum): | `SKIING` | `skiing.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/skiing.mp4) | """ # noqa: E501 // docs - VEHICLES = "vehicles.mp4" - MILK_BOTTLING_PLANT = "milk-bottling-plant.mp4" - VEHICLES_2 = "vehicles-2.mp4" - GROCERY_STORE = "grocery-store.mp4" - SUBWAY = "subway.mp4" - MARKET_SQUARE = "market-square.mp4" - PEOPLE_WALKING = "people-walking.mp4" - BEACH = "beach-1.mp4" - BASKETBALL = "basketball-1.mp4" - SKIING = "skiing.mp4" - - @classmethod - def list(cls) -> list[str]: - return list(map(lambda c: c.value, cls)) - - -VIDEO_ASSETS: dict[str, tuple[str, str]] = { - VideoAssets.VEHICLES.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.VEHICLES.value}", - "8155ff4e4de08cfa25f39de96483f918", - ), - VideoAssets.VEHICLES_2.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.VEHICLES_2.value}", - "830af6fba21ffbf14867a7fea595937b", - ), - VideoAssets.MILK_BOTTLING_PLANT.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.MILK_BOTTLING_PLANT.value}", + VEHICLES = ("vehicles.mp4", "8155ff4e4de08cfa25f39de96483f918") + MILK_BOTTLING_PLANT = ( + "milk-bottling-plant.mp4", "9e8fb6e883f842a38b3d34267290bdc7", - ), - VideoAssets.GROCERY_STORE.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.GROCERY_STORE.value}", - "11402e7b861c1980527d3d74cbe3b366", - ), - VideoAssets.SUBWAY.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.SUBWAY.value}", - "453475750691fb23c56a0cffef089194", - ), - VideoAssets.MARKET_SQUARE.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.MARKET_SQUARE.value}", - "859179bf4a21f80a8baabfdb2ed716dc", - ), - VideoAssets.PEOPLE_WALKING.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.PEOPLE_WALKING.value}", - "0574c053c8686c3f1dc0aa3743e45cb9", - ), - VideoAssets.BEACH.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.BEACH.value}", - "4175d42fec4d450ed081523fd39e0cf8", - ), - VideoAssets.BASKETBALL.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.BASKETBALL.value}", - "60d94a3c7c47d16f09d342b088012ecc", - ), - VideoAssets.SKIING.value: ( - f"{BASE_VIDEO_URL}{VideoAssets.SKIING.value}", - "d30987cbab1bbc5934199cdd1b293119", - ), + ) + VEHICLES_2 = ("vehicles-2.mp4", "830af6fba21ffbf14867a7fea595937b") + GROCERY_STORE = ("grocery-store.mp4", "48608fb4a8981f1c2469fa492adeec9c") + SUBWAY = ("subway.mp4", "453475750691fb23c56a0cffef089194") + MARKET_SQUARE = ("market-square.mp4", "859179bf4a21f80a8baabfdb2ed716dc") + PEOPLE_WALKING = ("people-walking.mp4", "0574c053c8686c3f1dc0aa3743e45cb9") + BEACH = ("beach-1.mp4", "4175d42fec4d450ed081523fd39e0cf8") + BASKETBALL = ("basketball-1.mp4", "60d94a3c7c47d16f09d342b088012ecc") + SKIING = ("skiing.mp4", "d30987cbab1bbc5934199cdd1b293119") + + +class ImageAssets(Assets): + """ + Each member of this enum represents a image asset. The value associated with each + member is the filename of the image. + + | Asset | Image Filename | Video URL | + |--------------------|------------------------|---------------------------------------------------------------------------------------| + | `PEOPLE_WALKING` | `people-walking.jpg` | [Link](https://media.roboflow.com/supervision/image-examples/people-walking.jpg) | + | `SOCCER` | `soccer.jpg` | [Link](https://media.roboflow.com/supervision/image-examples/soccer.jpg) | + + """ # noqa: E501 // docs + + PEOPLE_WALKING = ("people-walking.jpg", "e6bda00b47f2908eeae7df86ef995dcd") + SOCCER = ("soccer.jpg", "0f5a4b98abf3e3973faf9e9260a7d876") + + +MEDIA_ASSETS: dict[str, tuple[str, str]] = { + **{ + asset.filename: (f"{BASE_VIDEO_URL}{asset.filename}", asset.md5_hash) + for asset in VideoAssets + }, + **{ + asset.filename: (f"{BASE_IMAGE_URL}{asset.filename}", asset.md5_hash) + for asset in ImageAssets + }, } diff --git a/tests/assets/test_downloader.py b/tests/assets/test_downloader.py index 07549378..a4790bab 100644 --- a/tests/assets/test_downloader.py +++ b/tests/assets/test_downloader.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock, mock_open, patch import pytest from supervision.assets.downloader import download_assets, is_md5_hash_matching -from supervision.assets.list import VideoAssets +from supervision.assets.list import ImageAssets, VideoAssets class TestMD5HashMatching: @@ -51,13 +51,13 @@ class TestDownloadAssets: @patch("supervision.assets.downloader.is_md5_hash_matching", return_value=False) @patch("pathlib.Path.exists", return_value=True) def test_already_exists_but_corrupted( - self, mock_exists, mock_md5, mock_remove, mock_recursive + self, mock_exists, mock_md5, mock_remove, mock_download ): """Test download_assets when file exists but is corrupted (re-downloads).""" filename = "vehicles.mp4" result = download_assets(filename) assert result == filename - mock_recursive.assert_called_with(filename) + mock_download.assert_called_with(filename) @patch("builtins.print") @patch("pathlib.Path.open", new_callable=mock_open) @@ -108,6 +108,17 @@ class TestDownloadAssets: assert "Invalid asset" in str(exc_info.value) assert "vehicles.mp4" in str(exc_info.value) + @patch("pathlib.Path.exists", return_value=True) + def test_invalid_asset_when_file_exists(self, mock_exists): + """Test download_assets with invalid asset name that already exists.""" + invalid_filename = "invalid.mp4" + + with pytest.raises(ValueError, match="Invalid asset") as exc_info: + download_assets(invalid_filename) + + assert "Invalid asset" in str(exc_info.value) + assert "vehicles.mp4" in str(exc_info.value) + @patch("builtins.print") @patch("pathlib.Path.open", new_callable=mock_open) @patch("pathlib.Path.mkdir") @@ -115,7 +126,7 @@ class TestDownloadAssets: @patch("supervision.assets.downloader.tqdm") @patch("supervision.assets.downloader.get") @patch("pathlib.Path.exists", return_value=False) - def test_with_enum( + def test_with_video_enum( self, mock_exists, mock_get, @@ -138,4 +149,36 @@ class TestDownloadAssets: mock_tqdm.wrapattr.return_value.__exit__ = MagicMock() result = download_assets(asset) - assert result == asset.value + assert result == asset.filename + + @patch("builtins.print") + @patch("pathlib.Path.open", new_callable=mock_open) + @patch("pathlib.Path.mkdir") + @patch("supervision.assets.downloader.copyfileobj") + @patch("supervision.assets.downloader.tqdm") + @patch("supervision.assets.downloader.get") + @patch("pathlib.Path.exists", return_value=False) + def test_with_image_enum( + self, + mock_exists, + mock_get, + mock_tqdm, + mock_copyfileobj, + mock_mkdir, + mock_open_file, + mock_print, + ): + """Test download_assets with ImageAssets enum.""" + asset = ImageAssets.SOCCER + + mock_response = MagicMock() + mock_response.headers = {"Content-Length": "100"} + mock_response.raw = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_get.return_value = mock_response + + mock_tqdm.wrapattr.return_value.__enter__ = MagicMock() + mock_tqdm.wrapattr.return_value.__exit__ = MagicMock() + + result = download_assets(asset) + assert result == asset.filename diff --git a/tests/assets/test_list.py b/tests/assets/test_list.py index bf5dc47c..0447af28 100644 --- a/tests/assets/test_list.py +++ b/tests/assets/test_list.py @@ -1,4 +1,10 @@ -from supervision.assets.list import BASE_VIDEO_URL, VIDEO_ASSETS, VideoAssets +from supervision.assets.list import ( + BASE_IMAGE_URL, + BASE_VIDEO_URL, + MEDIA_ASSETS, + ImageAssets, + VideoAssets, +) def test_video_assets_list(): @@ -18,31 +24,45 @@ def test_video_assets_list(): assert VideoAssets.list() == expected_filenames -def test_video_assets_enum_values(): - """Test that VideoAssets enum members have correct values.""" +def test_image_assets_list(): + """Test that ImageAssets.list() returns all image filenames.""" + expected_filenames = [ + "people-walking.jpg", + "soccer.jpg", + ] + assert ImageAssets.list() == expected_filenames + + +def test_video_assets_values(): + """Test that VideoAssets enum members have correct attributes.""" + assert VideoAssets.VEHICLES.filename == "vehicles.mp4" + assert VideoAssets.VEHICLES.md5_hash == "8155ff4e4de08cfa25f39de96483f918" assert VideoAssets.VEHICLES.value == "vehicles.mp4" - assert VideoAssets.MILK_BOTTLING_PLANT.value == "milk-bottling-plant.mp4" - assert VideoAssets.VEHICLES_2.value == "vehicles-2.mp4" - assert VideoAssets.GROCERY_STORE.value == "grocery-store.mp4" - assert VideoAssets.SUBWAY.value == "subway.mp4" - assert VideoAssets.MARKET_SQUARE.value == "market-square.mp4" - assert VideoAssets.PEOPLE_WALKING.value == "people-walking.mp4" - assert VideoAssets.BEACH.value == "beach-1.mp4" - assert VideoAssets.BASKETBALL.value == "basketball-1.mp4" - assert VideoAssets.SKIING.value == "skiing.mp4" -def test_video_assets_dict_keys(): - """Test that VIDEO_ASSETS has all VideoAssets as keys.""" - expected_keys = {asset.value for asset in VideoAssets} - assert set(VIDEO_ASSETS.keys()) == expected_keys +def test_image_assets_values(): + """Test that ImageAssets enum members have correct attributes.""" + assert ImageAssets.SOCCER.filename == "soccer.jpg" + assert ImageAssets.SOCCER.md5_hash == "0f5a4b98abf3e3973faf9e9260a7d876" + assert ImageAssets.SOCCER.value == "soccer.jpg" -def test_video_assets_dict_values(): - """Test that VIDEO_ASSETS values are tuples of (url, md5_hash).""" - for filename, (url, md5_hash) in VIDEO_ASSETS.items(): +def test_media_assets_dict_keys(): + """Test that MEDIA_ASSETS has all VideoAssets and ImageAssets as keys.""" + expected_keys = {asset.filename for asset in VideoAssets} | { + asset.filename for asset in ImageAssets + } + assert set(MEDIA_ASSETS.keys()) == expected_keys + + +def test_media_assets_dict_values(): + """Test that MEDIA_ASSETS values are tuples of (url, md5_hash).""" + for filename, (url, md5_hash) in MEDIA_ASSETS.items(): assert isinstance(url, str) - assert url.startswith(BASE_VIDEO_URL) + if filename.endswith(".mp4"): + assert url.startswith(BASE_VIDEO_URL) + elif filename.endswith(".jpg"): + assert url.startswith(BASE_IMAGE_URL) assert url.endswith(filename) assert isinstance(md5_hash, str) assert len(md5_hash) == 32 # MD5 hash length