Add ImageAssets for download (#932)

* Add ImageAssets for download
* Update directory for consistency with video assets directory and add soccer image
* Fix unaligned table
* Add support for `ImageAssets` and extend `download_assets` functionality to handle both image and video assets.
* Refactor `Assets` enum initialization and improve `download_assets` logic with enhanced file checks and re-download handling. Add new tests for invalid asset cases and update existing test assertions.
* Simplify `download_assets` type handling and refactor `Assets` enum with additional attributes for filename and md5_hash.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: jirka <6035284+Borda@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 4 people

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.
 
 <div class="md-typeset">
@@ -18,3 +18,9 @@ that you can use in your demos.
 </div>
 
 :::supervision.assets.list.VideoAssets
+
+<div class="md-typeset">
+    <h2><a href="#supervision.assets.list.ImageAssets">ImageAssets</a></h2>
+</div>
+
+:::supervision.assets.list.ImageAssets

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"]

src/supervision/assets/downloader.py

@@ -8,7 +8,7 @@
 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
-
-    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()
-
-        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")
-
+    filename = asset_name.filename if isinstance(asset_name, Assets) else asset_name
+
+    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)
+
+            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}."
         )

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
+    },
 }

tests/assets/test_downloader.py

@@ -3,7 +3,7 @@
 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 @@ def test_already_exists_and_valid(self, mock_exists, mock_md5, mock_print):
     @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,14 +108,25 @@ def test_invalid_asset(self, mock_exists):
         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")
     @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_enum(
+    def test_with_video_enum(
         self,
         mock_exists,
         mock_get,
@@ -138,4 +149,36 @@ def test_with_enum(
         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