Merge pull request #6 from aws-deepracer-community/enhance-data-collection

Allow additional data collection

Author: jochem725

poetry.lock

@@ -11,7 +11,7 @@ boto3 = "^1.17.17"
 pandas = "^1.2.2"
 requests = "^2.25.1"
 joblib = "^1.0.1"
-deepracer-utils = "^0.19"
+deepracer-utils = "^0.20"
 
 [tool.poetry.dev-dependencies]
 pytest = "^5.2"

pyproject.toml

@@ -3,9 +3,13 @@
 # Track data.
 TRACK_FOLDER = "tracks"
 TRACK_FOLDER_ASSETS = os.path.join(TRACK_FOLDER, "assets")
+TRACK_FOLDER_ROUTES = os.path.join(TRACK_FOLDER, "npy")
 TRACK_CSV_FILEPATH = os.path.join(TRACK_FOLDER, "tracks.csv")
 
 # Leaderboard data
 LEADERBOARDS_FOLDER = "leaderboards"
 LEADERBOARDS_FOLDER_ASSETS = os.path.join(LEADERBOARDS_FOLDER, "assets")
 LEADERBOARDS_CSV_FILEPATH = os.path.join(LEADERBOARDS_FOLDER, "leaderboards.csv")
+
+# Simapp
+SIMAPP_TAR_GZ = "deepracer-simapp.tar.gz"

src/deepracer_race_stats/constants.py

@@ -1,14 +1,20 @@
 import os
 import click
+import boto3
+import glob
+import shutil
+import tarfile
 
 from datetime import datetime
 from joblib import Parallel, delayed
 from deepracer_race_stats.constants import (
     LEADERBOARDS_CSV_FILEPATH,
     LEADERBOARDS_FOLDER_ASSETS,
     LEADERBOARDS_FOLDER,
+    SIMAPP_TAR_GZ,
     TRACK_CSV_FILEPATH,
     TRACK_FOLDER_ASSETS,
+    TRACK_FOLDER_ROUTES,
 )
 
 from deepracer_race_stats.util.csv_util import boto_response_to_csv
@@ -17,7 +23,7 @@
     list_leaderboards,
     list_tracks,
 )
-from deepracer_race_stats.util.media import fetch_media_assets
+from deepracer_race_stats.util.assets import fetch_assets, extract_asset_paths
 
 
 @click.group()
@@ -43,11 +49,59 @@ def track_update(ctx, output_folder):
 
     boto_response_to_csv(response, output_path)
 
-    asset_map = {r["TrackArn"]: r["TrackPicture"] for r in response}
+    asset_map = {}
+
+    for r in response:
+        asset_map.update(extract_asset_paths(r, arn_key="TrackArn"))
 
     output_assets_folder = os.path.join(output_folder, TRACK_FOLDER_ASSETS)
+    fetch_assets(asset_map, output_assets_folder)
+
+
+@cli.command()
+@click.option("-o", "--output-folder", required=True)
+@click.option("-b", "--simapp-bucket", default="deepracer-managed-resources-us-east-1")
+@click.option("-k", "--simapp-key", default=SIMAPP_TAR_GZ)
+@click.pass_context
+def simapp_update(ctx, output_folder, simapp_bucket, simapp_key):
+    # Download simapp bundle.
+    tmp_folder = "simapp_tmp"
+    route_prefix = "opt/install/deepracer_simulation_environment/share/deepracer_simulation_environment/routes"
+
+    if not os.path.exists(tmp_folder):
+        os.makedirs(tmp_folder)
+
+    s3 = boto3.resource("s3")
+    s3.Bucket(simapp_bucket).download_file(Key=simapp_key, Filename=os.path.join("simapp_tmp", SIMAPP_TAR_GZ))
+
+    # Specify subfolders to extract.
+    def subfolders(tf):
+        for m in tf.getmembers():
+            # Routes folder with the track numpy files.
 
-    fetch_media_assets(asset_map, output_assets_folder)
+            if m.path.startswith(route_prefix):
+                yield m
+
+    with tarfile.open(os.path.join(tmp_folder, SIMAPP_TAR_GZ)) as f:
+        f.extractall(path=tmp_folder)
+
+    # Extract specific parts of the simapp we want to store.
+    with tarfile.open(os.path.join(tmp_folder, "bundle.tar"), mode="r") as f:
+        f.extractall(path=tmp_folder, members=subfolders(f))
+
+    # Move the files we want to the raw_data folder.
+    output_track_folder = os.path.join(output_folder, TRACK_FOLDER_ROUTES)
+
+    if not os.path.exists(output_track_folder):
+        os.makedirs(output_track_folder)
+
+    for route in glob.glob(os.path.join(tmp_folder, route_prefix, "*.npy")):
+        output_path = os.path.join(output_track_folder, os.path.basename(route))
+
+        shutil.copy(route, output_path)
+
+    # Remove temporary folder
+    shutil.rmtree(tmp_folder)
 
 
 @cli.command()
@@ -65,10 +119,12 @@ def leaderboard_update(ctx, output_folder):
 
     boto_response_to_csv(response, output_path)
 
-    asset_map = {r["Arn"]: r["ImageUrl"] for r in response if "ImageUrl" in r}
-    output_assets_folder = os.path.join(output_folder, LEADERBOARDS_FOLDER_ASSETS)
+    asset_map = {}
+    for r in response:
+        asset_map.update(extract_asset_paths(r))
 
-    fetch_media_assets(asset_map, output_assets_folder)
+    output_assets_folder = os.path.join(output_folder, LEADERBOARDS_FOLDER_ASSETS)
+    fetch_assets(asset_map, output_assets_folder)
 
     # Now do an update for each unique ARN:
     # - If OPEN: We collect a snapshot and save it under the current data and time.

src/deepracer_race_stats/main.py

@@ -0,0 +1,47 @@
+import os
+import requests
+from joblib import Parallel, delayed
+from urllib.parse import urlparse
+
+
+def fetch_assets(key_url_map, output_folder):
+    # Specific for tracks, we also collect the assets.
+    def download(key, url):
+        try:
+            output_path = os.path.join(output_folder, key)
+            output_dir = os.path.dirname(output_path)
+
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir, exist_ok=True)
+
+            r = requests.get(url)
+            if r.status_code == 200:
+                with open(output_path, "wb") as f:
+                    f.write(r.content)
+        except Exception as e:
+            print(e)
+            # Assume we can't get it then.
+            pass
+
+    return Parallel(n_jobs=-1, prefer="threads")(delayed(download)(key, url) for key, url in key_url_map.items())
+
+
+def get_asset_path(arn, url):
+    return os.path.join(arn, urlparse(url).path.lstrip("/"))
+
+
+def extract_asset_paths(r, arn_key="Arn"):
+    arn = r[arn_key]
+    response_asset_map = {}
+
+    columns = ["ImageUrl", "LeaderboardImage", "TrackPicture"]
+
+    for c in columns:
+        if c in r:
+            response_asset_map[get_asset_path(arn, r[c])] = r[c]
+
+    if "TrackRaceTypePictureMap" in r:
+        for key, value in r["TrackRaceTypePictureMap"].items():
+            response_asset_map[get_asset_path(arn, r["TrackRaceTypePictureMap"][key])] = value
+
+    return response_asset_map