|
import tarfile |
|
from dataclasses import dataclass |
|
from pathlib import Path |
|
from typing import List |
|
|
|
import boto3 |
|
from joblib import Parallel, delayed |
|
|
|
from realfake.utils import get_user_name, inject_args, Args |
|
|
|
|
|
class DownloadArgs(Args): |
|
start_idx: int = 0 |
|
end_idx: int = 5247 |
|
metadata_only: bool = False |
|
|
|
|
|
@inject_args |
|
def main(args: DownloadArgs) -> None: |
|
print(args) |
|
bucket, prefix = "s-datasets", "laion-aesthetic/data/laion2B-en-aesthetic/" |
|
start_idx, end_idx = args.start_idx, args.end_idx |
|
keys_range = list(range(start_idx, end_idx)) |
|
|
|
output_dir = Path(f"/fsx/{get_user_name()}/data/real_aes_{start_idx}_{end_idx}") |
|
if not args.metadata_only: |
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
metadata_dir = output_dir.parent/f"{output_dir.name}.metadata" |
|
metadata_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
jobs = get_jobs(keys_range, bucket, prefix, output_dir, metadata_dir, args.metadata_only) |
|
|
|
Parallel(n_jobs=-1, backend="multiprocessing", verbose=100)(delayed(download_and_extract)(job) for job in jobs) |
|
|
|
|
|
@dataclass |
|
class Job: |
|
bucket: str |
|
key: Path |
|
output_dir: Path |
|
|
|
|
|
def get_jobs( |
|
keys_range: list, |
|
bucket: str, |
|
prefix: str, |
|
output_dir: Path, |
|
metadata_dir: Path, |
|
metadata_only: bool, |
|
) -> List[Job]: |
|
client = boto3.client("s3") |
|
|
|
token, jobs = None, [] |
|
|
|
while True: |
|
conf = dict(Bucket=bucket, Prefix=prefix) |
|
if token is not None: conf["ContinuationToken"] = token |
|
response = client.list_objects_v2(**conf) |
|
|
|
for item in response.get("Contents"): |
|
key = Path(item["Key"]) |
|
if key.suffix == ".tar" and int(key.stem) in keys_range and not metadata_only: |
|
jobs.append(Job(bucket, key, output_dir)) |
|
elif key.suffix == ".parquet" and int(key.stem) in keys_range: |
|
jobs.append(Job(bucket, key, metadata_dir)) |
|
|
|
if not response["IsTruncated"]: break |
|
token = response["NextContinuationToken"] |
|
|
|
return jobs |
|
|
|
|
|
def download_and_extract(job: Job) -> None: |
|
client = boto3.client("s3") |
|
filename = job.output_dir / job.key.name |
|
|
|
print(f"{job.key}: downloading...") |
|
client.download_file(job.bucket, str(job.key), filename) |
|
|
|
if filename.suffix == ".tar": |
|
print(f"{job.key}: extracting...") |
|
with tarfile.open(filename) as tar: |
|
for name in tar.getnames(): |
|
extracted_path = job.output_dir/name |
|
if extracted_path.exists(): |
|
continue |
|
if name.endswith(".jpg"): |
|
tar.extract(name, job.output_dir) |
|
filename.unlink() |
|
|
|
print(f"{job.key}: done!") |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
