wehub-resource-sync
129 lines
6.0 KiB
Python
129 lines
6.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import aiohttp
|
|
import os
|
|
import shutil
|
|
from modelscope.hub.utils.utils import get_cache_dir
|
|
from typing import List, Literal, Optional, Union
|
|
|
|
from swift.utils import get_logger, safe_ddp_context
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class MediaResource:
|
|
"""A class to manage the resource downloading."""
|
|
|
|
cache_dir = os.path.join(get_cache_dir(), 'media_resources')
|
|
lock_dir = os.path.join(get_cache_dir(), 'lockers')
|
|
|
|
media_type_urls = {
|
|
'llava', 'coco', 'sam', 'gqa', 'ocr_vqa', 'textvqa', 'VG_100K', 'VG_100K_2', 'share_textvqa', 'web-celebrity',
|
|
'web-landmark', 'wikiart'
|
|
}
|
|
|
|
URL_PREFIX = 'https://www.modelscope.cn/api/v1/datasets/hjh0119/sharegpt4v-images/repo?Revision=master&FilePath='
|
|
|
|
@staticmethod
|
|
def get_url(media_type):
|
|
is_ocr_vqa = (media_type == 'ocr_vqa')
|
|
extension = 'tar' if is_ocr_vqa else 'zip'
|
|
return f'{MediaResource.URL_PREFIX}{media_type}.{extension}'
|
|
|
|
@staticmethod
|
|
def download(media_type_or_url: Union[str, List[str]],
|
|
local_alias: Optional[str] = None,
|
|
file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'):
|
|
"""Download and extract a resource from a http link.
|
|
|
|
Args:
|
|
media_type_or_url: `str` or List or `str`, Either belongs to the `media_type_urls`
|
|
listed in the class field, or a remote url to download and extract.
|
|
Be aware that, this media type or url needs to contain a zip or tar file.
|
|
local_alias: `Options[str]`, The local alias name for the `media_type_or_url`. If the first arg is a
|
|
media_type listed in this class, local_alias can leave None. else please pass in a name for the url.
|
|
The local dir contains the extracted files will be: {cache_dir}/{local_alias}
|
|
file_type: The file type, if is a compressed file, un-compressed the file,
|
|
if is an original file, only download it, if is a sharded file, download all files and extract.
|
|
|
|
Returns:
|
|
The local dir contains the extracted files.
|
|
"""
|
|
media_file = media_type_or_url if isinstance(media_type_or_url, str) else media_type_or_url[0]
|
|
with safe_ddp_context(hash_id=media_file):
|
|
return MediaResource._safe_download(
|
|
media_type=media_type_or_url, media_name=local_alias, file_type=file_type)
|
|
|
|
@staticmethod
|
|
def move_directory_contents(src_dir, dst_dir):
|
|
if not os.path.exists(dst_dir):
|
|
os.makedirs(dst_dir)
|
|
|
|
for dirpath, dirnames, filenames in os.walk(src_dir):
|
|
relative_path = os.path.relpath(dirpath, src_dir)
|
|
target_dir = os.path.join(dst_dir, relative_path)
|
|
|
|
if not os.path.exists(target_dir):
|
|
os.makedirs(target_dir)
|
|
|
|
for file in filenames:
|
|
src_file = os.path.join(dirpath, file)
|
|
dst_file = os.path.join(target_dir, file)
|
|
shutil.move(src_file, dst_file)
|
|
|
|
@staticmethod
|
|
def _safe_download(media_type: Union[str, List[str]],
|
|
media_name: Optional[str] = None,
|
|
file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'):
|
|
media_name = media_name or media_type
|
|
assert isinstance(media_name, str), f'{media_name} is not a str'
|
|
if isinstance(media_type, str) and media_type in MediaResource.media_type_urls:
|
|
media_type = MediaResource.get_url(media_type)
|
|
|
|
from datasets.download.download_manager import DownloadConfig, DownloadManager
|
|
final_folder = os.path.join(MediaResource.cache_dir, media_name)
|
|
|
|
if file_type == 'file':
|
|
filename = media_type.split('/')[-1]
|
|
final_path = os.path.join(final_folder, filename)
|
|
if os.path.exists(final_path): # if the download thing is a file but not folder,
|
|
return final_folder # check whether the file exists
|
|
if not os.path.exists(final_folder):
|
|
os.makedirs(final_folder) # and make sure final_folder exists to contain it
|
|
else:
|
|
if os.path.exists(final_folder):
|
|
return final_folder
|
|
|
|
logger.info('# #################Resource downloading#################')
|
|
logger.info('Downloading necessary resources...')
|
|
logger.info(f'Resource package: {media_type}')
|
|
logger.info(f'Extracting to local dir: {final_folder}')
|
|
logger.info('If the downloading fails or lasts a long time, '
|
|
'you can manually download the resources and extracting to the local dir.')
|
|
logger.info('Now begin.')
|
|
download_config = DownloadConfig(cache_dir=MediaResource.cache_dir)
|
|
download_config.storage_options = {'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=86400)}}
|
|
if file_type == 'file':
|
|
filename = media_type.split('/')[-1]
|
|
final_path = os.path.join(final_folder, filename)
|
|
local_dirs = DownloadManager(download_config=download_config).download(media_type)
|
|
shutil.move(str(local_dirs), final_path)
|
|
elif file_type == 'compressed':
|
|
local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_type)
|
|
shutil.move(str(local_dirs), final_folder)
|
|
else:
|
|
for media_url in media_type:
|
|
local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_url)
|
|
MediaResource.move_directory_contents(str(local_dirs), final_folder)
|
|
logger.info('# #################Resource downloading finished#################')
|
|
return final_folder
|
|
|
|
@staticmethod
|
|
def safe_save(image, file_name, folder, format='JPEG'):
|
|
folder = os.path.join(MediaResource.cache_dir, folder)
|
|
os.makedirs(folder, exist_ok=True)
|
|
file = os.path.join(folder, file_name)
|
|
if os.path.exists(file):
|
|
return file
|
|
image.save(file, format=format)
|
|
return file
|