import os import shutil from pathlib import Path from typing import Dict, List, TypedDict from zipfile import ZipFile import requests from PIL import Image from rich.progress import BarColumn, DownloadColumn, MofNCompleteColumn, Progress, TextColumn, TransferSpeedColumn from tqdm import tqdm """ { "name": "coco/train2017", # Visual Instruct Tuning images are all sourced from COCO Train 2017 "extract": True, "extract_type": "directory", "url": "http://images.cocodataset.org/zips/train2017.zip", "do_rename": True, }, { "name": "gqa/images", "extract": True, "extract_type": "directory", "url": "https://downloads.cs.stanford.edu/nlp/data/gqa/images.zip", "do_rename": True, }, { "name": "ocr_vqa/images", "extract": True, "extract_type": "directory", "url": "https://hf-mirror.com/datasets/qnguyen3/ocr_vqa/resolve/main/ocr_vqa.zip", "do_rename": True, }, { "name": "textvqa/train_images", "extract": True, "extract_type": "directory", "url": "https://dl.fbaipublicfiles.com/textvqa/images/train_val_images.zip", "do_rename": True, }, { "name": "vg/VG_100K_2", "extract": True, "extract_type": "directory", "url": "https://cs.stanford.edu/people/rak248/VG_100K_2/images2.zip", "do_rename": True, }, """ # === Dataset Registry w/ Links === # fmt: off DatasetComponent = TypedDict( "DatasetComponent", {"name": str, "extract": bool, "extract_type": str, "url": str, "do_rename": bool}, total=False ) DATASET_REGISTRY: Dict[str, List[DatasetComponent]] = { # === LLaVa v1.5 Dataset(s) === "llava-v1.5-instruct":[ { "name": "vg/VG_100K", "extract": True, "extract_type": "directory", "url": "https://cs.stanford.edu/people/rak248/VG_100K_2/images.zip", "do_rename": True, } ] } # fmt: on def convert_to_jpg(image_dir: Path) -> None: """Handling for OCR-VQA Images specifically; iterates through directory, converts all GIFs/PNGs.""" print(f"Converting all Images in `{image_dir}` to JPG") for image_fn in tqdm(list(image_dir.iterdir())): jpg_fn = image_dir / f"{image_fn.stem}.jpg" # 创建 JPG 文件名 if image_fn.suffix in {".jpg", ".jpeg"} or jpg_fn.exists(): continue if image_fn.suffix == ".gif": gif = Image.open(image_fn) gif.seek(0) gif.convert("RGB").save(jpg_fn) elif image_fn.suffix == ".png": Image.open(image_fn).convert("RGB").save(jpg_fn) else: raise ValueError(f"Unexpected image format `{image_fn.suffix}`") import os import shutil from pathlib import Path from typing import Dict, List, TypedDict from zipfile import ZipFile import requests from PIL import Image from rich.progress import BarColumn, DownloadColumn, MofNCompleteColumn, Progress, TextColumn, TransferSpeedColumn from tqdm import tqdm # DatasetComponent 和 DATASET_REGISTRY 保持不变 def download_with_progress(url: str, download_dir: Path, chunk_size_bytes: int = 1024) -> Path: """Utility function for downloading files from the internet, with a handy Rich-based progress bar.""" print(f"Downloading {url}") dest_path = download_dir / Path(url).name resume_header = {} if dest_path.exists(): return dest_path max_retries = 5 for attempt in range(max_retries): try: response = requests.get(url, headers=resume_header, stream=True) if response.status_code not in (200, 206): raise Exception(f"Failed to download. Status code: {response.status_code}") # 下载进度条 with Progress( TextColumn("[bold]{task.description} - {task.fields[fname]}"), BarColumn(bar_width=None), "[progress.percentage]{task.percentage:>3.1f}%", "•", DownloadColumn(), "•", TransferSpeedColumn(), transient=True, ) as dl_progress: dl_tid = dl_progress.add_task( "Downloading", fname=dest_path.name, total=int(response.headers.get("content-length", "None")) ) with open(dest_path, "ab") as f: # 以二进制追加模式打开文件 for data in response.iter_content(chunk_size=chunk_size_bytes): f.write(data) dl_progress.advance(dl_tid, chunk_size_bytes) return dest_path except Exception as e: print(f"Attempt {attempt + 1}/{max_retries} failed: {e}") if attempt < max_retries - 1: print("Retrying...") else: raise # 其他函数保持不变，main 方法也不变 def extract_with_progress(archive_path: Path, download_dir: Path, extract_type: str, cleanup: bool = False) -> Path: """Utility function for extracting compressed archives, with a handy Rich-based progress bar.""" assert archive_path.suffix == ".zip", "Only `.zip` compressed archives are supported for now!" print(f"Extracting {archive_path.name} to `{download_dir}`") with Progress( TextColumn("[bold]{task.description} - {task.fields[aname]}"), BarColumn(bar_width=None), "[progress.percentage]{task.percentage:>3.1f}%", "•", MofNCompleteColumn(), transient=True, ) as ext_progress: with ZipFile(archive_path) as zf: ext_tid = ext_progress.add_task("Extracting", aname=archive_path.name, total=len(members := zf.infolist())) extract_path = Path(zf.extract(members[0], download_dir)) if extract_type == "file": assert len(members) == 1, f"Archive `{archive_path}` with extract type `{extract_type} has > 1 member!" elif extract_type == "directory": for member in members[1:]: zf.extract(member, download_dir) ext_progress.advance(ext_tid) else: raise ValueError(f"Extract type `{extract_type}` for archive `{archive_path}` is not defined!") if cleanup: archive_path.unlink() return extract_path def download_extract(dataset_id: str, root_dir: Path) -> None: """Download all files for a given dataset (querying registry above), extracting archives if necessary.""" os.makedirs(download_dir := root_dir / "download" / dataset_id, exist_ok=True) # Download Files dl_tasks = [d for d in DATASET_REGISTRY[dataset_id] if not (download_dir / d["name"]).exists()] for dl_task in dl_tasks: dl_path = download_with_progress(dl_task["url"], download_dir) if dl_task["extract"]: dl_path = extract_with_progress(dl_path, download_dir, dl_task["extract_type"]) dl_path = dl_path.parent if dl_path.is_file() else dl_path if dl_task["do_rename"]: shutil.move(dl_path, download_dir / dl_task["name"]) if __name__ == "__main__": import sys from pathlib import Path # 设置根目录 root_dir = Path("./data") # 这里设置一个默认的下载目录 os.makedirs(root_dir, exist_ok=True) # 下载所有数据集 for dataset_id in DATASET_REGISTRY.keys(): print(f"开始下载数据集: {dataset_id}") download_extract(dataset_id, root_dir) print("所有数据集下载完成！")