From 537c50e45f94b214338c6e53bc9822b15fe3a595 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 7 May 2024 08:58:08 +0200 Subject: [PATCH] Fix NCNN export on Windows (#11719) Signed-off-by: Glenn Jocher --- ultralytics/engine/exporter.py | 15 ++++++++------- ultralytics/utils/checks.py | 2 +- ultralytics/utils/downloads.py | 2 +- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 5ab8ed29d6..8a073f59c3 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -88,7 +88,7 @@ from ultralytics.utils import ( yaml_save, ) from ultralytics.utils.checks import check_imgsz, check_is_path_safe, check_requirements, check_version -from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, is_url +from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download from ultralytics.utils.files import file_size, spaces_in_path from ultralytics.utils.ops import Profile from ultralytics.utils.torch_utils import TORCH_1_13, get_latest_opset, select_device, smart_inference_mode @@ -543,19 +543,20 @@ class Exporter: system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux" try: _, assets = get_github_assets(repo="pnnx/pnnx") - url = [x for x in assets if f"{system}.zip" in x][0] - assert is_url(url), "Unable to retrieve PNNX repo assets" - LOGGER.info(f"Successfully retrieve PNNX repo asset URL {url}") + asset = [x for x in assets if f"{system}.zip" in x][0] + assert isinstance(asset, str), "Unable to retrieve PNNX repo assets" # i.e. pnnx-20240410-macos.zip + LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}") + asset = attempt_download_asset(asset, repo="pnnx/pnnx", release="latest") except Exception as e: url = f"https://github.com/pnnx/pnnx/releases/download/20240410/pnnx-20240410-{system}.zip" LOGGER.warning(f"{prefix} WARNING ⚠️ PNNX GitHub assets not found: {e}, using default {url}") - asset = attempt_download_asset(url, repo="pnnx/pnnx", release="latest") + asset = safe_download(url) if check_is_path_safe(Path.cwd(), asset): # avoid path traversal security vulnerability unzip_dir = Path(asset).with_suffix("") (unzip_dir / name).rename(pnnx) # move binary to ROOT - shutil.rmtree(unzip_dir) # delete unzip dir - Path(asset).unlink() # delete zip pnnx.chmod(0o777) # set read, write, and execute permissions for everyone + shutil.rmtree(unzip_dir) # delete unzip dir + Path(asset).unlink(missing_ok=True) # delete zip ncnn_args = [ f'ncnnparam={f / "model.ncnn.param"}', diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py index 65365ef450..e5f5c80a22 100644 --- a/ultralytics/utils/checks.py +++ b/ultralytics/utils/checks.py @@ -528,7 +528,7 @@ def check_is_path_safe(basedir, path): base_dir_resolved = Path(basedir).resolve() path_resolved = Path(path).resolve() - return path_resolved.is_file() and path_resolved.parts[: len(base_dir_resolved.parts)] == base_dir_resolved.parts + return path_resolved.exists() and path_resolved.parts[: len(base_dir_resolved.parts)] == base_dir_resolved.parts def check_imshow(warn=False): diff --git a/ultralytics/utils/downloads.py b/ultralytics/utils/downloads.py index 9f4b0e0a40..eb979d2a1d 100644 --- a/ultralytics/utils/downloads.py +++ b/ultralytics/utils/downloads.py @@ -418,7 +418,7 @@ def attempt_download_asset(file, repo="ultralytics/assets", release="v8.2.0", ** Example: ```python - file_path = attempt_download_asset('yolov5s.pt', repo='ultralytics/assets', release='latest') + file_path = attempt_download_asset('yolov8n.pt', repo='ultralytics/assets', release='latest') ``` """ from ultralytics.utils import SETTINGS # scoped for circular import