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>
2026-02-12 19:22:15 +09:00
1 parent acbaf4ff25
commit 0ed48f0227
@@ -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
@@ -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"]
@@ -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}."
        )
@@ -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
+    },
 }
@@ -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
@@ -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