Add ImageAssets for download

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.
 
 ## Install extra
@@ -29,3 +29,9 @@ as an extra within the Supervision package.
 </div>
 
 :::supervision.assets.list.VideoAssets
+
+<div class="md-typeset">
+  <h2>ImageAssets</h2>
+</div>
+
+:::supervision.assets.list.ImageAssets

supervision/assets/__init__.py

@@ -1,2 +1,2 @@
 from supervision.assets.downloader import download_assets
-from supervision.assets.list import VideoAssets
+from supervision.assets.list import ImageAssets, VideoAssets

supervision/assets/downloader.py

@@ -4,7 +4,7 @@
 from shutil import copyfileobj
 from typing import Union
 
-from supervision.assets.list import VIDEO_ASSETS, VideoAssets
+from supervision.assets.list import ASSETS, Assets, ImageAssets, VideoAssets
 
 try:
     from requests import get
@@ -41,31 +41,34 @@ 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: Union[VideoAssets, str]) -> str:
+def download_assets(asset_name: Union[VideoAssets, ImageAssets, str]) -> str:
     """
     Download a specified asset if it doesn't already exist or is corrupted.
 
     Parameters:
-        asset_name (Union[VideoAssets, str]): The name or type of the asset to be
+        asset_name (Union[VideoAssets, ImageAssets, str]): The name or type of the asset to be
             downloaded.
 
     Returns:
         str: The filename of the downloaded asset.
 
     Example:
         ```python
-        from supervision.assets import download_assets, VideoAssets
+        from supervision.assets import download_assets, VideoAssets, ImageAssets
 
         download_assets(VideoAssets.VEHICLES)
         "vehicles.mp4"
+
+        download_assets(ImageAssets.PEOPLE_WALKING)
+        "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:
+    if not Path(filename).exists() and filename in ASSETS:
         print(f"Downloading {filename} assets \n")
-        response = get(VIDEO_ASSETS[filename][0], stream=True, allow_redirects=True)
+        response = get(ASSETS[filename][0], stream=True, allow_redirects=True)
         response.raise_for_status()
 
         file_size = int(response.headers.get("Content-Length", 0))
@@ -79,15 +82,15 @@ def download_assets(asset_name: Union[VideoAssets, str]) -> str:
                 copyfileobj(raw_resp, file)
 
     elif Path(filename).exists():
-        if not is_md5_hash_matching(filename, VIDEO_ASSETS[filename][1]):
+        if not is_md5_hash_matching(filename, 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 ASSETS.keys())
         raise ValueError(
             f"Invalid asset. It should be one of the following: {valid_assets}."
         )

supervision/assets/list.py

@@ -2,14 +2,21 @@
 from typing import Dict, Tuple
 
 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):
+    def __init__(self, filename: str, hash: str):
+        self.filename = filename
+        self.hash = hash
+
+
+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) |
@@ -20,46 +27,41 @@ class VideoAssets(Enum):
     | `PEOPLE_WALKING`       | `people-walking.mp4`       | [Link](https://media.roboflow.com/supervision/video-examples/people-walking.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"
+    VEHICLES = ("vehicles.mp4", "8155ff4e4de08cfa25f39de96483f918")
+    MILK_BOTTLING_PLANT = (
+        "milk-bottling-plant.mp4",
+        "9e8fb6e883f842a38b3d34267290bdc7",
+    )
+    VEHICLES_2 = ("vehicles-2.mp4", "830af6fba21ffbf14867a7fea595937b")
+    GROCERY_STORE = ("grocery-store.mp4", "453475750691fb23c56a0cffef089194")
+    SUBWAY = ("subway.mp4", "453475750691fb23c56a0cffef089194")
+    MARKET_SQUARE = ("market-square.mp4", "859179bf4a21f80a8baabfdb2ed716dc")
+    PEOPLE_WALKING = ("people-walking.mp4", "0574c053c8686c3f1dc0aa3743e45cb9")
 
-    @classmethod
-    def list(cls):
-        return list(map(lambda c: c.value, cls))
 
+class ImageAssets(Assets):
+    """
+    Each member of this enum represents a image asset. The value associated with each
+    member is the filename of the image.
 
-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}",
-        "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",
-    ),
+    | 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")
+
+
+ASSETS: Dict[str, Tuple[str, str]] = {
+    **{
+        asset.value[0]: (f"{BASE_VIDEO_URL}{asset.value[0]}", asset.value[1])
+        for asset in VideoAssets
+    },
+    **{
+        asset.value[0]: (f"{BASE_IMAGE_URL}{asset.value[0]}", asset.value[1])
+        for asset in ImageAssets
+    },
 }