data_download.py

"""
download.py

Utility functions for downloading and extracting various datasets to (local) disk.
"""

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

#from prismatic.overwatch import initialize_overwatch

# Initialize Overwatch =>> Wraps `logging.Logger`
#overwatch = initialize_overwatch(__name__)


# === 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) ===

    # Note =>> This is the full suite of datasets included in the LLaVa 1.5 "finetuning" stage; all the LLaVa v1.5
    #          models are finetuned on this split. We use this dataset for all experiments in our paper.
    "llava-v1.5-instruct":
    [
        {
            "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",
            "extract": True,
            "extract_type": "directory",
            "url": "https://cs.stanford.edu/people/rak248/VG_100K_2/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,
        },
    ]
}
# 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())):
        if image_fn.suffix in {".jpg", ".jpeg"} or (jpg_fn := image_dir / f"{image_fn.stem}.jpg").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}`")


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 {(dest_path := download_dir / Path(url).name)} from `{url}`", ctx_level=1)
    if dest_path.exists():
        return dest_path

    # Otherwise --> fire an HTTP Request, with `stream = True`
    response = requests.get(url, stream=True)

    # Download w/ Transfer-Aware Progress
    #   => Reference: https://github.com/Textualize/rich/blob/master/examples/downloader.py
    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, "wb") as f:
            for data in response.iter_content(chunk_size=chunk_size_bytes):
                dl_progress.advance(dl_tid, f.write(data))

    return dest_path


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}`", ctx_level=1)

    # Extract w/ Progress
    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!")

    # Cleanup (if specified)
    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 => Single-Threaded, with Progress Bar
    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)

        # Extract Files (if specified) --> Note (assumes ".zip" ONLY!)
        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

        # Rename Path --> dl_task["name"]
        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("所有数据集下载完成！")