Skip to content

Commit

Permalink
Fix the pull_asset code; Update minor docstring (#516)
Browse files Browse the repository at this point in the history
* Fix the asset remove code; Update minor docstring

* fix bug
  • Loading branch information
pengzhenghao committed Oct 19, 2023
1 parent ee8447e commit aa3446d
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 10 deletions.
14 changes: 9 additions & 5 deletions metadrive/pull_asset.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import argparse
import os
import progressbar
import shutil
import urllib.request
import zipfile

import progressbar
import filelock

from metadrive.constants import VERSION
from metadrive.engine.logger import get_logger
Expand Down Expand Up @@ -55,17 +55,21 @@ def pull_asset(update):
return

zip_path = os.path.join(TARGET_DIR, 'assets.zip')
zip_lock = os.path.join(TARGET_DIR, 'assets.zip.lock')

# Fetch the zip file
logger.info("Pull assets from {}".format(ASSET_URL))
urllib.request.urlretrieve(ASSET_URL, zip_path, MyProgressBar())

# Extract the zip file to the desired location
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(TARGET_DIR)
lock = filelock.FileLock(zip_lock)
with lock:
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(TARGET_DIR)

# Remove the downloaded zip file (optional)
os.remove(zip_path)
if os.path.exists(zip_path):
os.remove(zip_path)
logger.info("Successfully download assets, version: {}. MetaDrive version: {}".format(asset_version(), VERSION))


Expand Down
25 changes: 21 additions & 4 deletions metadrive/scenario/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.pyplot import figure

from metadrive.component.static_object.traffic_object import TrafficCone, TrafficBarrier
from metadrive.component.traffic_light.base_traffic_light import BaseTrafficLight
from metadrive.component.traffic_participants.cyclist import Cyclist
Expand Down Expand Up @@ -323,6 +322,14 @@ def convert_recorded_scenario_exported(record_episode, scenario_log_interval=0.1


def read_scenario_data(file_path):
"""Read a scenario pkl file and return the Scenario Description instance.
Args:
file_path: the path to a scenario file (usually ends with `.pkl`).
Returns:
The Scenario Description instance of that scenario.
"""
assert SD.is_scenario_file(file_path), "File: {} is not scenario file".format(file_path)
with open(file_path, "rb") as f:
# unpickler = CustomUnpickler(f)
Expand All @@ -333,11 +340,21 @@ def read_scenario_data(file_path):

def read_dataset_summary(file_folder, check_file_existence=True):
"""
We now support two methods to load pickle files.
Read the `dataset_summary.pkl` and return the metadata of each scenario in this dataset.
This function supports two methods to load pickle files. The first is the old method where we store pickle files in
0.pkl, 1.pkl, .... The second is the new method which use a summary file to record important metadata of
each scenario.
The first is the old method where we store pickle files in 0.pkl, 1.pkl, ...
Args:
file_folder: the path to the root folder of your dataset.
check_file_existence: check if all scenarios registered in the summary file exist.
The second is the new method which use a summary file to record important metadata of each scenario.
Returns:
A tuple of three elements:
1) the summary dict mapping from scenario ID to its metadata,
2) the list of all scenarios IDs, and
3) a dict mapping from scenario IDs to the folder that hosts their files.
"""
summary_file = os.path.join(file_folder, SD.DATASET.SUMMARY_FILE)
mapping_file = os.path.join(file_folder, SD.DATASET.MAPPING_FILE)
Expand Down
3 changes: 2 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,8 @@ def is_win():
"scipy",
"psutil",
"geopandas",
"shapely"
"shapely",
"filelock"
]

nuplan_requirement = [
Expand Down

0 comments on commit aa3446d

Please sign in to comment.