From 575055ce7028ee99618ff1c5c0e8919e8e2cd849 Mon Sep 17 00:00:00 2001 From: Kalen Michael Date: Fri, 28 Oct 2022 21:16:03 +0200 Subject: [PATCH] Switch from suffix checks to archive checks (#9963) * fix: switched from suffix checks to archive checks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup seems like both functions accept Path type input so removing str() * import always Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- utils/general.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/utils/general.py b/utils/general.py index 46978f1..88cefb7 100644 --- a/utils/general.py +++ b/utils/general.py @@ -23,8 +23,9 @@ from itertools import repeat from multiprocessing.pool import ThreadPool from pathlib import Path from subprocess import check_output +from tarfile import is_tarfile from typing import Optional -from zipfile import ZipFile +from zipfile import ZipFile, is_zipfile import cv2 import IPython @@ -465,7 +466,7 @@ def check_dataset(data, autodownload=True): # Download (optional) extract_dir = '' - if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip + if isinstance(data, (str, Path)) and (is_zipfile(data) or is_tarfile(data)): download(data, dir=f'{DATASETS_DIR}/{Path(data).stem}', unzip=True, delete=False, curl=False, threads=1) data = next((DATASETS_DIR / Path(data).stem).rglob('*.yaml')) extract_dir, autodownload = data.parent, False @@ -607,11 +608,11 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1, retry else: LOGGER.warning(f'❌ Failed to download {url}...') - if unzip and success and f.suffix in ('.zip', '.tar', '.gz'): + if unzip and success and (f.suffix == '.gz' or is_zipfile(f) or is_tarfile(f)): LOGGER.info(f'Unzipping {f}...') - if f.suffix == '.zip': + if is_zipfile(f): unzip_file(f, dir) # unzip - elif f.suffix == '.tar': + elif is_tarfile(f): os.system(f'tar xf {f} --directory {f.parent}') # unzip elif f.suffix == '.gz': os.system(f'tar xfz {f} --directory {f.parent}') # unzip