Dataset Downloader (#5)

Author: bbengfort

.gitignore

@@ -161,4 +161,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-datasets
+construe/datasets/fixtures

MANIFEST.in

@@ -3,8 +3,8 @@ include *.rst
 include *.txt
 include *.yml
 include *.cfg
-include MANIFEST.in
 
+include MANIFEST.in
 include LICENSE
 
 graft docs
@@ -14,6 +14,7 @@ graft tests
 prune tests/fixtures
 
 graft construe
+prune construe/datasets/fixtures
 
 global-exclude __pycache__
 global-exclude *.py[co]

construe/__main__.py

@@ -6,10 +6,12 @@
 import torch
 import platform
 
-from click import ClickException
 from .version import get_version
 from .basic import BasicBenchmark
+from .exceptions import DeviceError
 from .moondream import MoonDreamBenchmark
+from .datasets.manifest import generate_manifest
+from .datasets.path import get_data_home, FIXTURES
 
 
 CONTEXT_SETTINGS = {
@@ -34,13 +36,26 @@
     envvar=["CONSTRUE_ENV", "ENV"],
     help="name of the experimental environment for comparison (default is hostname)",
 )
+@click.option(
+    "-D",
+    "--datadir",
+    default=None,
+    envvar="CONSTRUE_DATA",
+    help="specify the location to download datasets to",
+)
+@click.option(
+    "-C",
+    "--cleanup/--no-cleanup",
+    default=True,
+    help="cleanup all downloaded datasets after the benchmark is run",
+)
 @click.pass_context
-def main(ctx, env=None, device=None):
+def main(ctx, env=None, device=None, datadir=None, cleanup=True):
     if device is not None:
         try:
             torch.set_default_device(device)
         except RuntimeError as e:
-            raise ClickException(str(e))
+            raise DeviceError(e)
 
         click.echo(f"using torch.device(\"{device}\")")
 
@@ -50,6 +65,8 @@ def main(ctx, env=None, device=None):
     ctx.ensure_object(dict)
     ctx.obj["device"] = device
     ctx.obj["env"] = env
+    ctx.obj["data_home"] = get_data_home(datadir)
+    ctx.obj["cleanup"] = cleanup
 
 
 @main.command()
@@ -89,11 +106,30 @@ def basic(ctx, **kwargs):
 @main.command()
 @click.pass_context
 def moondream(ctx, **kwargs):
-    kwargs["env"] = ctx["env"]
+    kwargs["env"] = ctx.obj["env"]
     benchmark = MoonDreamBenchmark(**kwargs)
     benchmark.run()
 
 
+@main.command()
+@click.option(
+    "-f",
+    "--fixtures",
+    type=str,
+    default=FIXTURES,
+    help="path to fixtures directory to generate manifest from",
+)
+@click.option(
+    "-o",
+    "--out",
+    type=str,
+    default=None,
+    help="path to write the manifest to",
+)
+def manifest(fixtures=FIXTURES, out=None):
+    generate_manifest(fixtures, out)
+
+
 if __name__ == "__main__":
     main(
         obj={},

construe/datasets.py

@@ -0,0 +1,7 @@
+"""
+Manages datasets used for inferencing
+"""
+
+from .loaders import * # noqa
+from .download import download_data
+from .path import get_data_home, cleanup_dataset

construe/datasets/__init__.py

@@ -0,0 +1,68 @@
+"""
+Handle downloading datasets from our content URL
+"""
+
+import os
+import zipfile
+
+from urllib.request import urlopen
+
+from .signature import sha256sum
+from .path import get_data_home, cleanup_dataset
+
+from construe.exceptions import DatasetsError
+
+
+# Downlod chunk size
+CHUNK = 524288
+
+
+def download_data(url, signature, data_home=None, replace=False, extract=True):
+    """
+    Downloads the zipped data set specified at the given URL, saving it to
+    the data directory specified by ``get_data_home``. This function verifies
+    the download with the given signature and extracts the archive.
+    """
+    data_home = get_data_home(data_home)
+
+    # Get the name of the file from the URL
+    basename = os.path.basename(url)
+    name, _ = os.path.splitext(basename)
+
+    # Get the archive and data directory paths
+    archive = os.path.join(data_home, basename)
+    datadir = os.path.join(data_home, name)
+
+    # If the archive exists cleanup or raise override exception
+    if os.path.exists(archive):
+        if not replace:
+            raise DatasetsError(
+                ("dataset already exists at {}, set replace=False to overwrite").format(
+                    archive
+                )
+            )
+
+        cleanup_dataset(name, data_home=data_home)
+
+    # Create the output directory if it does not exist
+    if not os.path.exists(datadir):
+        os.mkdir(datadir)
+
+    # Fetch the response in a streaming fashion and write it to disk.
+    response = urlopen(url)
+
+    with open(archive, "wb") as f:
+        while True:
+            chunk = response.read(CHUNK)
+            if not chunk:
+                break
+            f.write(chunk)
+
+    # Compare the signature of the archive to the expected one
+    if sha256sum(archive) != signature:
+        raise ValueError("Download signature does not match hardcoded signature!")
+
+    # If extract, extract the zipfile.
+    if extract:
+        zf = zipfile.ZipFile(archive)
+        zf.extractall(path=data_home)

construe/datasets/download.py

@@ -0,0 +1,47 @@
+"""
+Managers for loading datasets
+"""
+
+import os
+import glob
+
+from .manifest import load_manifest
+from .download import download_data
+from ..exceptions import DatasetsError
+from .path import dataset_archive, find_dataset_path, cleanup_dataset
+
+
+__all__ = ["load_content_moderation", "cleanup_content_moderation"]
+
+
+DATASETS = load_manifest()
+CONTENT_MODERATION = "content-moderation"
+
+
+def _info(dataset):
+    if dataset not in DATASETS:
+        raise DatasetsError(f"no dataset named {dataset} exists")
+    return DATASETS[dataset]
+
+
+def load_content_moderation(data_home=None):
+    """
+    Downloads the content moderation dataset if it does not exist then
+    yields all of the paths for the images in the dataset.
+    """
+    info = _info(CONTENT_MODERATION)
+    if not dataset_archive(CONTENT_MODERATION, info["signature"], data_home=data_home):
+        # If the dataset does not exist, download and extract it
+        info.update({"data_home": data_home, "replace": False, "extract": True})
+        download_data(**info)
+
+    data_path = find_dataset_path(CONTENT_MODERATION, fname=None, ext=None)
+    for path in glob.glob(os.path.join(data_path, "**", "*")):
+        yield path
+
+
+def cleanup_content_moderation(data_home=None):
+    """
+    Removes the content moderation dataset and archive.
+    """
+    return cleanup_dataset(CONTENT_MODERATION, data_home=data_home)

construe/datasets/loaders.py

@@ -0,0 +1,6 @@
+{
+  "content-moderation": {
+    "url": "https://storage.googleapis.com/construe/v0.2.0/content-moderation.zip",
+    "signature": "4925f2733bf3b1596e4d950e49d12bf4fff2379f08cbf0a105807ff03d4c2b4e"
+  }
+}

construe/datasets/manifest.json

@@ -0,0 +1,40 @@
+"""
+Manifest handlers for downloading and checking signatures
+"""
+
+import os
+import json
+import glob
+
+from urllib.parse import urljoin
+
+from .signature import sha256sum
+from ..version import get_version
+from .path import FIXTURES, MANIFEST
+
+
+BASE_URL = "https://storage.googleapis.com/"
+
+
+def load_manifest(path=MANIFEST):
+    with open(MANIFEST, "r") as f:
+        return json.load(f)
+
+
+def generate_manifest(fixtures=FIXTURES, out=MANIFEST):
+    out = out or MANIFEST
+
+    manifest = {}
+    version = get_version(short=True)
+
+    for path in glob.glob(os.path.join(fixtures, "*.zip")):
+        fname = os.path.basename(path)
+        name, _ = os.path.splitext(fname)
+
+        manifest[name] = {
+            "url": urljoin(BASE_URL, f"construe/v{version}/{fname}"),
+            "signature": sha256sum(path),
+        }
+
+    with open(out, "w") as o:
+        json.dump(manifest, o, indent=2)