diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..2e1cedc --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,33 @@ +name: CI + +on: + push: + branches: [dev, main] + pull_request: + branches: [dev, main] + +jobs: + lint: + name: ruff (lint) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - run: pip install ruff==0.15.21 + - run: ruff check src tests + + typecheck: + name: mypy (advisory) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - run: pip install mypy + # Advisory: the codebase is partially typed. Print findings but never + # fail the build (green check) until type coverage improves. Tighten by + # dropping the `|| true` once the baseline is cleaned up. + - run: mypy src/gasbench --ignore-missing-imports || true diff --git a/pyproject.toml b/pyproject.toml index 7b877c4..2f1842a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "gasbench" -version = "0.7.5" +version = "0.8.0" description = "GASBench - ML model benchmark evaluation package" readme = "README.md" requires-python = ">=3.10" @@ -41,6 +41,13 @@ dependencies = [ "open-clip-torch==3.3.0", ] +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "ruff==0.15.21", + "mypy>=1.10", +] + [project.scripts] gasbench = "gasbench.cli:main" @@ -57,6 +64,23 @@ where = ["src"] line-length = 88 target-version = ['py310'] +[tool.ruff] +line-length = 88 +target-version = "py310" +src = ["src", "tests"] + +[tool.ruff.lint] +# High-signal baseline: pyflakes (real bugs) + import/statement errors. +# Style rules (line length, naming) are intentionally left off for now to keep +# the diff focused; expand this list over time. +select = ["E4", "E7", "E9", "F"] + +[tool.mypy] +python_version = "3.10" +ignore_missing_imports = true +follow_imports = "silent" +# Advisory only for now ā€” the codebase is partially typed. Not enforced in CI. + [tool.pytest.ini_options] testpaths = ["tests"] markers = ["slow: marks tests as slow (network I/O)"] diff --git a/src/gasbench/benchmark.py b/src/gasbench/benchmark.py index 60c41f3..26ab657 100644 --- a/src/gasbench/benchmark.py +++ b/src/gasbench/benchmark.py @@ -3,7 +3,6 @@ import json import time import uuid -import asyncio from datetime import datetime from pathlib import Path from typing import Dict, Optional @@ -181,7 +180,7 @@ async def load_model_for_benchmark( benchmark_results["validation"]["input_type"] = str(input_specs[0].type) benchmark_results["validation"]["output_shape"] = str(output_specs[0].shape) - logger.info(f"Model loaded successfully") + logger.info("Model loaded successfully") # Get execution providers providers = session.get_providers() @@ -237,7 +236,7 @@ async def execute_benchmark( if dataset_filters: logger.info(f"Dataset filters: {dataset_filters}") if modality == "image": - df = await run_image_benchmark( + await run_image_benchmark( session, input_specs, benchmark_results, @@ -264,7 +263,7 @@ async def execute_benchmark( ) benchmark_score = benchmark_results.get("image_results", {}).get("benchmark_score", 0.0) elif modality == "video": - df = await run_video_benchmark( + await run_video_benchmark( session, input_specs, benchmark_results, @@ -291,7 +290,7 @@ async def execute_benchmark( ) benchmark_score = benchmark_results.get("video_results", {}).get("benchmark_score", 0.0) elif modality == "audio": - df = await run_audio_benchmark( + await run_audio_benchmark( session, input_specs, benchmark_results, @@ -375,7 +374,7 @@ def format_benchmark_summary(benchmark_results: Dict) -> str: per_dataset = results.get("per_dataset_results", {}) if per_dataset: - lines.append(f"\nšŸ“‹ PER-DATASET RESULTS:") + lines.append("\nšŸ“‹ PER-DATASET RESULTS:") for dataset_name, dataset_results in per_dataset.items(): accuracy = dataset_results.get("accuracy", 0.0) total = dataset_results.get("total", 0) @@ -394,7 +393,7 @@ def format_benchmark_summary(benchmark_results: Dict) -> str: entry["samples"] += int(ds_stats.get("total", 0)) entry["correct"] += int(ds_stats.get("correct", 0)) if by_type: - lines.append(f"\nšŸŽ­ ACCURACY BY MEDIA TYPE:") + lines.append("\nšŸŽ­ ACCURACY BY MEDIA TYPE:") for mtype, v in by_type.items(): samples = v["samples"] correct = v["correct"] @@ -403,7 +402,7 @@ def format_benchmark_summary(benchmark_results: Dict) -> str: validation = benchmark_results.get("validation", {}) if validation: - lines.append(f"\nšŸ” MODEL VALIDATION:") + lines.append("\nšŸ” MODEL VALIDATION:") lines.append(f" Input Shape: {validation.get('input_shape', 'N/A')}") lines.append(f" Input Type: {validation.get('input_type', 'N/A')}") lines.append(f" Output Shape: {validation.get('output_shape', 'N/A')}") diff --git a/src/gasbench/benchmarks/audio_bench.py b/src/gasbench/benchmarks/audio_bench.py index 114a706..40a639f 100644 --- a/src/gasbench/benchmarks/audio_bench.py +++ b/src/gasbench/benchmarks/audio_bench.py @@ -1,6 +1,4 @@ import os -import time -import traceback import numpy as np from typing import Dict, Optional @@ -8,9 +6,14 @@ from ..processing.media import process_audio_sample from ..dataset.iterator import DatasetIterator -from .utils.inference import process_model_output from .recording import BenchmarkRunRecorder, log_dataset_summary -from .common import BenchmarkRunConfig, build_plan, create_tracker, finalize_run +from .common import ( + BenchmarkRunConfig, + build_plan, + create_tracker, + finalize_run, + run_batch_and_record, +) import pandas as pd logger = get_logger(__name__) @@ -40,49 +43,22 @@ def process_batch( # Squeeze any extra dimensions (e.g., (1, 96000) -> (96000,)) batch_audio_np = [b.squeeze() if b.ndim > 1 else b for b in batch_audio_np] - + batch_array = np.stack(batch_audio_np) except Exception as e: logger.error(f"Failed to stack audio batch: {e}") - return - - start = time.time() - outputs = None - try: - outputs = session.run(None, {input_specs[0].name: batch_array}) - except Exception as e: - logger.error(f"Inference failed: {e} (batch shape: {batch_array.shape})") - for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( - batch_metadata - ): + for label, sample, sample_index, dataset_name, sample_seed in batch_metadata: tracker.add_error( dataset_name=dataset_name, sample_index=sample_index, sample=sample, - error_message=f"inference-failed: {str(e)[:160]}", + error_message=f"stack-failed: {str(e)[:160]}", ) return - batch_inference_time = (time.time() - start) * 1000 - per_sample_time = batch_inference_time / len(batch_audio) - - for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( - batch_metadata - ): - predicted, pred_probs = process_model_output(outputs[0][i]) - - tracker.add_ok( - dataset_name=dataset_name, - sample_index=sample_index, - sample=sample, - label=label, - predicted=predicted, - probs=pred_probs, - inference_time_ms=per_sample_time, - batch_inference_time_ms=batch_inference_time, - batch_id=batch_id, - batch_size=len(batch_audio), - sample_seed=sample_seed, - ) + + run_batch_and_record( + session, input_specs, batch_array, batch_metadata, tracker, batch_id + ) async def run_audio_benchmark( diff --git a/src/gasbench/benchmarks/common.py b/src/gasbench/benchmarks/common.py index 25e409d..0d7570a 100644 --- a/src/gasbench/benchmarks/common.py +++ b/src/gasbench/benchmarks/common.py @@ -1,7 +1,10 @@ import json +import time from dataclasses import dataclass from typing import Dict, List, Optional, Tuple +import numpy as np + from ..config import DEFAULT_TARGET_SIZE from ..dataset.config import ( get_benchmark_size, @@ -19,9 +22,75 @@ compute_generator_stats_from_df, ) from .utils import calculate_per_source_accuracy +from .utils.inference import process_model_output from ..logger import get_logger +def stack_uniform_batch(items: List[np.ndarray]) -> np.ndarray: + """Stack a list of same-shaped arrays into a single batched array.""" + first = items[0] + batch_array = np.empty((len(items),) + first.shape, dtype=first.dtype) + for i, item in enumerate(items): + batch_array[i] = item + return batch_array + + +def run_batch_and_record( + session, + input_specs, + batch_array: np.ndarray, + batch_metadata, + tracker: BenchmarkRunRecorder, + batch_id: int, + aug_pass: bool = False, +): + """Run one batch through the model and record rows in the tracker. + + On inference failure every sample in the batch is recorded as an error + (rather than propagating and aborting the whole dataset), so a model that + crashes on hard samples cannot silently drop them from its score. + """ + logger = get_logger(__name__) + if not batch_metadata: + return + + start = time.time() + try: + outputs = session.run(None, {input_specs[0].name: batch_array}) + except Exception as e: + logger.error(f"Inference failed: {e} (batch shape: {batch_array.shape})") + for label, sample, sample_index, dataset_name, sample_seed in batch_metadata: + tracker.add_error( + dataset_name=dataset_name, + sample_index=sample_index, + sample=sample, + error_message=f"inference-failed: {str(e)[:160]}", + ) + return + + batch_inference_time = (time.time() - start) * 1000 + per_sample_time = batch_inference_time / len(batch_metadata) + + for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( + batch_metadata + ): + predicted, pred_probs = process_model_output(outputs[0][i]) + tracker.add_ok( + dataset_name=dataset_name, + sample_index=sample_index, + sample=sample, + label=label, + predicted=predicted, + probs=pred_probs, + inference_time_ms=per_sample_time, + batch_inference_time_ms=batch_inference_time, + batch_id=batch_id, + batch_size=len(batch_metadata), + sample_seed=sample_seed, + aug_pass=aug_pass, + ) + + @dataclass class BenchmarkRunConfig: modality: str diff --git a/src/gasbench/benchmarks/image_bench.py b/src/gasbench/benchmarks/image_bench.py index 8e8b0ed..fbe05e5 100644 --- a/src/gasbench/benchmarks/image_bench.py +++ b/src/gasbench/benchmarks/image_bench.py @@ -1,5 +1,4 @@ import os -import time import traceback import numpy as np from typing import Dict, Optional @@ -19,9 +18,15 @@ from ..dataset.iterator import DatasetIterator -from .utils.inference import process_model_output from .recording import BenchmarkRunRecorder, log_dataset_summary, build_sample_id -from .common import BenchmarkRunConfig, build_plan, create_tracker, finalize_run +from .common import ( + BenchmarkRunConfig, + build_plan, + create_tracker, + finalize_run, + stack_uniform_batch, + run_batch_and_record, +) from .aug_cache import img_aug_cache_path, write_aug_cache import pandas as pd @@ -219,36 +224,10 @@ def process_batch( """push a batch of images through the model and record rows in tracker.""" if not batch_images: return - - first = batch_images[0] - batch_array = np.empty((len(batch_images),) + first.shape, dtype=first.dtype) - for i, img in enumerate(batch_images): - batch_array[i] = img - - start = time.time() - outputs = session.run(None, {input_specs[0].name: batch_array}) - batch_inference_time = (time.time() - start) * 1000 - per_sample_time = batch_inference_time / len(batch_images) - - for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( - batch_metadata - ): - predicted, pred_probs = process_model_output(outputs[0][i]) - - tracker.add_ok( - dataset_name=dataset_name, - sample_index=sample_index, - sample=sample, - label=label, - predicted=predicted, - probs=pred_probs, - inference_time_ms=per_sample_time, - batch_inference_time_ms=batch_inference_time, - batch_id=batch_id, - batch_size=len(batch_images), - sample_seed=sample_seed, - aug_pass=aug_pass, - ) + batch_array = stack_uniform_batch(batch_images) + run_batch_and_record( + session, input_specs, batch_array, batch_metadata, tracker, batch_id, aug_pass + ) async def run_image_benchmark( diff --git a/src/gasbench/benchmarks/utils/inference.py b/src/gasbench/benchmarks/utils/inference.py index 9f6e64f..9658ffa 100644 --- a/src/gasbench/benchmarks/utils/inference.py +++ b/src/gasbench/benchmarks/utils/inference.py @@ -96,15 +96,22 @@ def process_model_output(logits: np.ndarray) -> Tuple[int, np.ndarray]: - probabilities: softmax probabilities as 1D array (length 1,2, or 3) """ logits = np.atleast_1d(logits).flatten() + + if len(logits) == 1: + p = 1.0 / (1.0 + np.exp(-logits[0])) + pred_probs = np.array([p], dtype=np.float64) + predicted_binary = int(p > 0.5) + return predicted_binary, pred_probs + exp_x = np.exp(logits - np.max(logits)) pred_probs = exp_x / np.sum(exp_x) - + if len(pred_probs) == 3: predicted_binary = 0 if np.argmax(pred_probs) == 0 else 1 elif len(pred_probs) == 2: predicted_binary = int(np.argmax(pred_probs)) else: predicted_binary = int(pred_probs[0] > 0.5) - + return predicted_binary, pred_probs diff --git a/src/gasbench/benchmarks/utils/metrics.py b/src/gasbench/benchmarks/utils/metrics.py index 245d801..b9d8e55 100644 --- a/src/gasbench/benchmarks/utils/metrics.py +++ b/src/gasbench/benchmarks/utils/metrics.py @@ -117,7 +117,6 @@ def compute_sn34_score( Returns: float in [0,1]. """ - import math # Normalize MCC: [-1, 1] -> [0, 1] bin_mcc = float(self.calculate_binary_mcc()) diff --git a/src/gasbench/benchmarks/utils/pytorch_session.py b/src/gasbench/benchmarks/utils/pytorch_session.py index 21020f1..55a0d5f 100644 --- a/src/gasbench/benchmarks/utils/pytorch_session.py +++ b/src/gasbench/benchmarks/utils/pytorch_session.py @@ -4,9 +4,9 @@ import torch import numpy as np from pathlib import Path -from typing import List, Any, Optional +from typing import List, Any -from .custom_model_loader import load_custom_model, load_model_config +from .custom_model_loader import load_custom_model from ...logger import get_logger logger = get_logger(__name__) @@ -156,13 +156,37 @@ def run(self, output_names: Any, input_dict: dict) -> List[np.ndarray]: with torch.no_grad(): output = self.model(tensor) - # Handle various output formats if hasattr(output, "logits"): - output = output.logits # HuggingFace models - elif isinstance(output, tuple): - output = output[0] # Some models return tuples + output = output.logits elif isinstance(output, dict): - # Some models return dicts with 'logits' key - output = output.get("logits", output.get("output", list(output.values())[0])) + if "logits" in output: + output = output["logits"] + elif "output" in output: + output = output["output"] + elif len(output) == 1: + output = next(iter(output.values())) + else: + raise ValueError( + f"Model returned a dict without a 'logits' or 'output' key; " + f"cannot determine which of {sorted(output.keys())} is the logits tensor. " + f"Return a logits tensor or key it as 'logits'." + ) + elif isinstance(output, (tuple, list)): + if len(output) == 1: + output = output[0] + else: + raise ValueError( + f"Model returned a {type(output).__name__} of length {len(output)}; " + f"cannot determine which element is the logits tensor. " + f"Return a logits tensor rather than a {type(output).__name__}." + ) + + last_dim = output.shape[-1] + if last_dim not in (1, 2, 3) and last_dim != self._num_classes: + raise ValueError( + f"Model output last dimension is {last_dim}, expected 1, 2, 3, " + f"or num_classes ({self._num_classes}). Check that model.py returns " + f"logits and that model_config.yaml declares the correct num_classes." + ) return [output.cpu().to(torch.float32).numpy()] diff --git a/src/gasbench/benchmarks/video_bench.py b/src/gasbench/benchmarks/video_bench.py index 6a8f490..f32bd5b 100644 --- a/src/gasbench/benchmarks/video_bench.py +++ b/src/gasbench/benchmarks/video_bench.py @@ -1,6 +1,4 @@ import os -import time -import asyncio import numpy as np import traceback from typing import Dict, Optional @@ -20,9 +18,15 @@ from ..constants import MAX_VIDEO_NUM_FRAMES from ..dataset.iterator import DatasetIterator -from .utils.inference import process_model_output from .recording import BenchmarkRunRecorder, log_dataset_summary, build_sample_id -from .common import BenchmarkRunConfig, build_plan, create_tracker, finalize_run +from .common import ( + BenchmarkRunConfig, + build_plan, + create_tracker, + finalize_run, + stack_uniform_batch, + run_batch_and_record, +) from .aug_cache import vid_aug_cache_path, write_aug_cache import pandas as pd @@ -239,48 +243,10 @@ def process_video_batch( """push a batch of videos through the model and record rows in tracker.""" if not batch_videos: return - - first = batch_videos[0] - batch_array = np.empty((len(batch_videos),) + first.shape, dtype=first.dtype) - for i, vid in enumerate(batch_videos): - batch_array[i] = vid - - start = time.time() - outputs = None - try: - outputs = session.run(None, {input_specs[0].name: batch_array}) - except Exception as e: - for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( - batch_metadata - ): - tracker.add_error( - dataset_name=dataset_name, - sample_index=sample_index, - sample=sample, - error_message=f"inference-failed: {str(e)[:160]}", - ) - return - batch_inference_time = (time.time() - start) * 1000 - per_sample_time = batch_inference_time / len(batch_videos) - - for i, (label, sample, sample_index, dataset_name, sample_seed) in enumerate( - batch_metadata - ): - predicted, pred_probs = process_model_output(outputs[0][i]) - tracker.add_ok( - dataset_name=dataset_name, - sample_index=sample_index, - sample=sample, - label=label, - predicted=predicted, - probs=pred_probs, - inference_time_ms=per_sample_time, - batch_inference_time_ms=batch_inference_time, - batch_id=batch_id, - batch_size=len(batch_videos), - sample_seed=sample_seed, - aug_pass=aug_pass, - ) + batch_array = stack_uniform_batch(batch_videos) + run_batch_and_record( + session, input_specs, batch_array, batch_metadata, tracker, batch_id, aug_pass + ) async def run_video_benchmark( @@ -385,8 +351,6 @@ async def run_video_benchmark( f"{dataset_config.name} ({dataset_cap} samples)" ) - dataset_skipped = 0 - try: is_gasstation = "gasstation" in dataset_config.name.lower() if skip_missing: diff --git a/src/gasbench/cli.py b/src/gasbench/cli.py index 79f24dd..7ac6184 100644 --- a/src/gasbench/cli.py +++ b/src/gasbench/cli.py @@ -101,15 +101,15 @@ def command_run(args): print(f"šŸ“¦ Using custom PyTorch model from: {model_path}") elif model_path.suffix.lower() == ".onnx": # ONNX format no longer supported for competition - print(f"Error: ONNX format is no longer supported for competition.") - print(f"Please use a safetensors model directory containing:") - print(f" - model_config.yaml") - print(f" - model.py") - print(f" - *.safetensors weights") - print(f"See docs/Safetensors.md for requirements.") + print("Error: ONNX format is no longer supported for competition.") + print("Please use a safetensors model directory containing:") + print(" - model_config.yaml") + print(" - model.py") + print(" - *.safetensors weights") + print("See docs/Safetensors.md for requirements.") return 1 else: - print(f"Error: Model must be a directory with model_config.yaml, model.py, and *.safetensors") + print("Error: Model must be a directory with model_config.yaml, model.py, and *.safetensors") print(f"Got: {model_path}") return 1 @@ -182,9 +182,9 @@ def command_run(args): f.write(format_benchmark_summary(results)) print(f"\nšŸ“Š Results saved to: {results_path}") - print(f" - JSON: results.json") - print(f" - Parquet: records.parquet") - print(f" - Summary: summary.txt") + print(" - JSON: results.json") + print(" - Parquet: records.parquet") + print(" - Summary: summary.txt") if results.get("benchmark_completed"): print("\nāœ… Benchmark completed successfully") @@ -207,7 +207,6 @@ def command_verify_cache(args): from .dataset.cache import ( scan_cache_directory, compute_cache_statistics, - verify_cache_against_configs, format_size_bytes, ) @@ -230,19 +229,19 @@ def command_verify_cache(args): by_modality, by_media_type = compute_cache_statistics(datasets) - print(f"\nšŸ“Š CACHE SUMMARY") + print("\nšŸ“Š CACHE SUMMARY") print(f"Total Datasets: {len(datasets)}") print(f"Total Samples: {total_samples:,}") print(f"Total Size: {format_size_bytes(total_size_bytes)}") - print(f"\nšŸ“ˆ BY MODALITY") + print("\nšŸ“ˆ BY MODALITY") for modality in sorted(by_modality.keys()): stats = by_modality[modality] print( f" {modality.upper():12} {stats['count']:3} datasets, {stats['samples']:7,} samples, {format_size_bytes(stats['size']):>10}" ) - print(f"\nšŸŽ­ BY MEDIA TYPE") + print("\nšŸŽ­ BY MEDIA TYPE") for mtype in sorted(by_media_type.keys()): stats = by_media_type[mtype] print( @@ -250,7 +249,7 @@ def command_verify_cache(args): ) if args.verbose: - print(f"\nšŸ“‹ DETAILED LISTING") + print("\nšŸ“‹ DETAILED LISTING") datasets_by_mod = defaultdict(list) for ds in datasets: @@ -304,7 +303,7 @@ def _verify_against_config(cached_names, cached_datasets, args, cache_dir): if missing: has_issues = True - print(f"\nāŒ MISSING DATASETS (in config but not in cache):") + print("\nāŒ MISSING DATASETS (in config but not in cache):") by_modality = defaultdict(list) for ds in expected_datasets: if ds.name in missing: @@ -318,7 +317,7 @@ def _verify_against_config(cached_names, cached_datasets, args, cache_dir): if extra: has_issues = True - print(f"\nāš ļø EXTRA DATASETS (in cache but not in config):") + print("\nāš ļø EXTRA DATASETS (in cache but not in config):") by_modality = defaultdict(list) for ds in cached_datasets: if ds["name"] in extra: diff --git a/src/gasbench/dataset/cache.py b/src/gasbench/dataset/cache.py index 0a7d370..d7bbc13 100644 --- a/src/gasbench/dataset/cache.py +++ b/src/gasbench/dataset/cache.py @@ -32,7 +32,7 @@ def check_dataset_cache( ): try: with open(dataset_info_file, "r") as f: - dataset_info = json.load(f) + json.load(f) # validate JSON is readable sample_count = len( [ diff --git a/src/gasbench/dataset/config.py b/src/gasbench/dataset/config.py index 650e837..d53197a 100644 --- a/src/gasbench/dataset/config.py +++ b/src/gasbench/dataset/config.py @@ -2,10 +2,8 @@ from typing import Dict, List, Optional, Tuple from pathlib import Path import os -import json import yaml from importlib.resources import files -from collections import defaultdict import hashlib from ..logger import get_logger @@ -404,7 +402,7 @@ def load_datasets_from_yaml(yaml_path: str) -> Dict[str, List[BenchmarkDatasetCo raise yaml.YAMLError(f"Failed to parse YAML file {yaml_path}: {e}") if not isinstance(data, dict): - raise ValueError(f"YAML file must contain a dictionary at root level") + raise ValueError("YAML file must contain a dictionary at root level") all_errors = [] result = {"image": [], "video": [], "audio": []} diff --git a/src/gasbench/dataset/download.py b/src/gasbench/dataset/download.py deleted file mode 100644 index f3070d6..0000000 --- a/src/gasbench/dataset/download.py +++ /dev/null @@ -1,1883 +0,0 @@ -import base64 -import hashlib -import io -import json -import os -import random -import re -import shutil -import tarfile -import tempfile -import time -import traceback -from urllib.parse import quote -from contextlib import closing -from io import BytesIO -from pathlib import Path -from typing import Any, Dict, Generator, List, Optional, Tuple -from zipfile import ZipFile -from datetime import datetime - -import soundfile as sf - -import pandas as pd -import pyarrow.parquet as pq -import requests -import huggingface_hub as hf_hub -from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError -import numpy as np -from PIL import Image -from modelscope.hub.api import HubApi as MSHubApi - -from ..logger import get_logger -from .utils.s3_utils import list_s3_files, download_s3_file, _get_s3_urls -from .utils.metadata_utils import create_sample, extract_row_metadata -from .utils.gasstation_utils import ( - extract_iso_week_from_path, - filter_files_by_current_week, - filter_files_by_recent_weeks, -) - - -class DatasetAccessError(Exception): - """Raised when a dataset cannot be accessed (gated, not found, or permission denied).""" - - pass - -logger = get_logger(__name__) - - -IMAGE_FILE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"} -VIDEO_FILE_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".webm", ".m4v", ".mpeg", ".mpg"} -AUDIO_FILE_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac"} - - -def _calculate_files_to_download( - dataset, - source_format: str, - media_per_archive: int, - archives_per_dataset: int, -) -> int: - """Calculate # files to download based on dataset modality and source format. - Returns -1 to indicate "download all files", or a positive integer for the count. - """ - src_fmt = source_format.lower().lstrip(".") - - # Direct media files (jpg, png, mp4, etc.) - download equivalent to archive extraction - if dataset.modality == "image": - is_direct_media = src_fmt in {ext.lstrip(".") for ext in IMAGE_FILE_EXTENSIONS} - elif dataset.modality == "audio": - is_direct_media = src_fmt in {ext.lstrip(".") for ext in AUDIO_FILE_EXTENSIONS} - else: - is_direct_media = src_fmt in {ext.lstrip(".") for ext in VIDEO_FILE_EXTENSIONS} - - if is_direct_media: - if media_per_archive == -1 or archives_per_dataset == -1: - return -1 - return media_per_archive * archives_per_dataset - - # Archive-based media (parquet, zip, tar, etc.) - return archives_per_dataset - - -def download_and_extract( - dataset, # BenchmarkDatasetConfig - media_per_archive: int = 100, - archives_per_dataset: int = 5, - temp_dir: Optional[str] = None, - force_download: bool = False, - current_week_only: bool = False, - num_weeks: int = None, - cache_dir: str = "/.cache/gasbench", - downloaded_archives: Optional[set] = None, - target_week: Optional[str] = None, - hf_token: Optional[str] = None, - seed: Optional[int] = None, -) -> Generator[Dict[str, Any], None, None]: - """ - Download datasets and yield extracted media as a generator. - - Downloads files temporarily, extracts/processing content, then cleans up downloads. - Processed content (videos, images) is cached persistently in other locations. - - Args: - dataset: BenchmarkDatasetConfig object - media_per_archive: Number of media items (images/videos) to extract per archive file - archives_per_dataset: Number of archive files to download per dataset - temp_dir: Temporary directory for downloads - force_download: Force download even if dataset appears to be cached - cache_dir: Base cache directory for persistent storage - - Yields: - Complete sample dictionaries ready for processing functions - """ - try: - if not force_download and _is_dataset_cached(dataset, cache_dir): - logger.info(f"Dataset {dataset.name} found in cache, loading from volume") - try: - cached_samples_count = 0 - for sample in _load_dataset_from_cache(dataset, cache_dir): - cached_samples_count += 1 - yield sample - logger.info( - f"Successfully loaded {cached_samples_count} samples from cache for {dataset.name}" - ) - return - except (OSError, json.JSONDecodeError, FileNotFoundError) as e: - logger.warning(f"Failed to load {dataset.name} from cache: {e}") - logger.info( - f"Fallback to download: {dataset.name} will be downloaded fresh" - ) - # Fall through to download logic - - # Create temporary directory for downloads - if temp_dir is not None: - try: - Path(temp_dir).mkdir(parents=True, exist_ok=True) - except OSError: - pass - temp_dir_root = Path(tempfile.mkdtemp(dir=temp_dir)) - - try: - include_paths = getattr(dataset, "include_paths", None) - exclude_paths = getattr(dataset, "exclude_paths", None) - source = getattr(dataset, "source", "huggingface") - - n_files = _calculate_files_to_download( - dataset, - dataset.source_format, - media_per_archive, - archives_per_dataset, - ) - max_files_to_list = n_files if n_files > 0 else None - - # ā”€ā”€ Filtered sequential mode ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ - # When filter_column/filter_value are set, download parquet shards - # one at a time and stop as soon as the sample target is hit. - filter_column = getattr(dataset, "filter_column", None) - filter_value = getattr(dataset, "filter_value", None) - if filter_column and filter_value: - try: - all_filenames = _list_remote_dataset_files( - dataset.path, - dataset.source_format or "parquet", - False, None, None, - include_paths, exclude_paths, - source, hf_token, - max_files=None, # need full list to shuffle and pick from - ) - except DatasetAccessError as e: - logger.warning(f"Skipping {dataset.name}: {e}") - return - - if not all_filenames: - logger.warning(f"No parquet files found for {dataset.path}") - return - - target_total = media_per_archive * archives_per_dataset - logger.info( - f"Filtered mode: {dataset.name} ā€” {filter_column}=={filter_value}, " - f"target {target_total} samples from {len(all_filenames)} shards" - ) - yield from _download_filtered_sequential( - dataset=dataset, - filenames=all_filenames, - target_total=target_total, - temp_dir_root=temp_dir_root, - source=source, - hf_token=hf_token, - seed=seed, - ) - return - # ā”€ā”€ Standard (non-filtered) download path ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ - - try: - filenames = _list_remote_dataset_files( - dataset.path, - dataset.source_format, - current_week_only, - num_weeks, - target_week, - include_paths, - exclude_paths, - source, - hf_token, - max_files=max_files_to_list, - ) - except DatasetAccessError as e: - logger.warning( - f"Skipping dataset {dataset.name}: {e}" - ) - return - - if not filenames: - is_gasstation = "gasstation" in dataset.name.lower() - if is_gasstation: - logger.warning( - f"No files found for {dataset.path} with format {dataset.source_format}. " - f"Gasstation datasets require parquet metadata files." - ) - return - - logger.warning( - f"No files found for {dataset.path} with format {dataset.source_format}" - ) - - src_fmt = str(getattr(dataset, "source_format", "")).lower().lstrip(".") - fallback_formats = [".parquet", ".zip", ".tar", ".tar.gz"] - for fallback_format in fallback_formats: - if fallback_format != dataset.source_format: - logger.info( - f"Trying fallback format {fallback_format} for {dataset.path}" - ) - try: - filenames = _list_remote_dataset_files( - dataset.path, - fallback_format, - current_week_only, - num_weeks, - target_week, - include_paths, - exclude_paths, - source, - hf_token, - max_files=max_files_to_list, - ) - except DatasetAccessError as e: - logger.warning( - f"Skipping dataset {dataset.name}: {e}" - ) - return - if filenames: - logger.info( - f"Found {len(filenames)} files with format {fallback_format}" - ) - # Recalculate n_files for the actual format found ā€” the - # original n_files was computed for source_format (e.g. mp4) - # which treats each file as one media item, but an archive - # format (parquet/zip/tar) should only download - # archives_per_dataset files, not media_per_archive. - n_files = _calculate_files_to_download( - dataset, fallback_format, media_per_archive, archives_per_dataset - ) - break - - if not filenames: - logger.warning(f"No files found for {dataset.path} with any format") - return - - remote_paths = _get_download_urls(dataset.path, filenames, source) - - is_gasstation = "gasstation" in dataset.name.lower() - to_download = _select_files_to_download( - remote_paths, n_files, prioritize_recent=is_gasstation, seed=seed - ) - - if is_gasstation and downloaded_archives is not None: - original_count = len(to_download) - to_download = [ - url - for url in to_download - if os.path.basename(url) not in downloaded_archives - ] - filtered_count = original_count - len(to_download) - if filtered_count > 0: - logger.info( - f"Skipping {filtered_count} already-downloaded archives, " - f"downloading {len(to_download)} new archives" - ) - - if len(to_download) == 0: - return - - logger.info( - f"Downloading {len(to_download)} files from {dataset.path} (dataset: {dataset.name})" - ) - - if ( - is_gasstation - and to_download - and any(".parquet" in url for url in to_download) - ): - yield from _process_gasstation( - dataset=dataset, - to_download=to_download, - temp_dir_root=temp_dir_root, - media_per_archive=media_per_archive, - downloaded_archives=downloaded_archives, - source=source, - hf_token=hf_token, - seed=seed, - ) - else: - max_download_workers = min(10, len(to_download)) - - successfully_processed = 0 - any_downloaded = False - for downloaded_file in _stream_downloads( - to_download, - temp_dir_root, - chunk_size=8192, - max_workers=max_download_workers, - hf_token=hf_token, - ): - if not downloaded_file or not downloaded_file.exists(): - continue - - any_downloaded = True - try: - iso_week = extract_iso_week_from_path(str(downloaded_file)) - successfully_processed += 1 - - for sample in yield_media_from_source( - downloaded_file, - dataset, - media_per_archive, - iso_week, - hf_token, - seed, - ): - yield sample - finally: - try: - if downloaded_file.exists(): - downloaded_file.unlink() - except OSError: - pass - - if not any_downloaded: - logger.warning( - f"No files successfully downloaded for {dataset.name}" - ) - return - - logger.info( - f"Downloaded and processed {successfully_processed} files " - f"for {dataset.name}" - ) - - finally: - if temp_dir_root.exists(): - shutil.rmtree(temp_dir_root) - - except Exception as e: - logger.error(f"Error processing {dataset.path}: {e}") - - -def yield_media_from_source( - source_path: Path, - dataset, # BenchmarkDatasetConfig - num_items: int, - iso_week: Optional[str] = None, - hf_token: Optional[str] = None, - seed: Optional[int] = None, -) -> Generator[Dict[str, Any], None, None]: - """ - Unified media extractor for parquet, zip, tar sources, frame directories, and raw media files. - - Returns complete sample dictionaries ready for processing functions. - Samples include: image/video_bytes, media_type, dataset_name, dataset_path, etc. - """ - try: - if source_path.is_dir(): - yield from _process_frame_directory(source_path, dataset, iso_week) - return - - filename = str(source_path.name).lower() - - if _is_parquet_file(filename): - yield from _process_parquet(source_path, dataset, num_items, iso_week, seed) - return - - if _is_zip_file(filename) or _is_tar_file(filename): - yield from _process_zip_or_tar( - source_path, - dataset, - num_items, - iso_week, - seed, - ) - return - - if any( - filename.endswith(ext) - for ext in (IMAGE_FILE_EXTENSIONS | VIDEO_FILE_EXTENSIONS | AUDIO_FILE_EXTENSIONS) - ): - yield from _process_raw(source_path, dataset, iso_week) - return - - logger.warning(f"Unsupported source format for {source_path}") - return - except Exception as e: - logger.warning(f"Error in yield_media_from_source for {source_path}: {e}") - return - - -def _select_files_to_download( - urls: List[str], - count: int, - prioritize_recent: bool = False, - seed: Optional[int] = None, -) -> List[str]: - """Select files to download. - - Args: - urls: List of file URLs to select from - count: Number of files to select (-1 = all files) - prioritize_recent: If True, sort by filename (descending) to get newest first. - Useful for gasstation datasets to get freshest data. - seed: Random seed for reproducible sampling (non-gasstation only) - - Returns: - List of selected URLs - """ - if count == -1: - return urls - if count <= 0: - return [] - - if prioritize_recent: - sorted_urls = sorted(urls, reverse=True) - return sorted_urls[: min(count, len(urls))] - - if seed is not None: - rng = random.Random(seed) - return rng.sample(urls, min(count, len(urls))) - return random.sample(urls, min(count, len(urls))) - - -MAX_FILES_DEFAULT = 5000 - - -def _list_remote_dataset_files( - dataset_path: str, - source_format: str = ".parquet", - current_week_only: bool = False, - num_weeks: int = None, - target_week: str = None, - include_paths: Optional[List[str]] = None, - exclude_paths: Optional[List[str]] = None, - source: str = "huggingface", - hf_token: Optional[str] = None, - max_files: Optional[int] = None, -) -> List[str]: - """List available files in a dataset, filtered by source_format and path patterns. - - Supports single extensions (e.g., .parquet, .zip) and tar variants (.tar, .tar.gz, .tgz). - For gasstation datasets with current_week_only=True, filters to only current ISO week's data. - For gasstation datasets with num_weeks set, filters to only the N most recent weeks. - - Args: - dataset_path: Dataset repository path (org/dataset-name) - source_format: File extension to filter by - current_week_only: For gasstation datasets, only get current week - num_weeks: For gasstation datasets, get N most recent weeks - target_week: For gasstation datasets, get specific week - include_paths: Only include files containing one of these path segments - exclude_paths: Exclude files containing any of these path segments - source: Source platform ("huggingface" or "modelscope") - max_files: Stop after collecting this many files (None uses MAX_FILES_DEFAULT) - """ - if max_files is None: - max_files = MAX_FILES_DEFAULT - # Don't add "." prefix for special formats like "frames" - if source_format != "frames" and not source_format.startswith("."): - source_format = "." + source_format - - if source_format in [".tar", ".tar.gz", ".tgz"]: - source_format = [".tar", ".tar.gz", ".tgz"] - - # For gasstation datasets with week filtering, we must fetch ALL files before - # applying max_files. Otherwise early termination picks up files from older - # weeks (alphabetically first) and the subsequent week filter discards them all, - # resulting in zero files. - is_gasstation = "gasstation" in dataset_path.lower() - needs_week_filter = is_gasstation and (target_week or num_weeks or current_week_only) - listing_max = None if needs_week_filter else max_files - - if source == "modelscope": - files = list_modelscope_files(repo_id=dataset_path, extension=source_format) - if include_paths: - files = [f for f in files if any(path_seg in f for path_seg in include_paths)] - if exclude_paths: - files = [f for f in files if not any(path_seg in f for path_seg in exclude_paths)] - if not needs_week_filter: - files = files[:max_files] - elif source == "s3": - files = list_s3_files(path=dataset_path, extension=source_format) - if include_paths: - files = [f for f in files if any(path_seg in f for path_seg in include_paths)] - if exclude_paths: - files = [f for f in files if not any(path_seg in f for path_seg in exclude_paths)] - if not needs_week_filter: - files = files[:max_files] - else: # hf - supports early termination natively - files = list_hf_files( - repo_id=dataset_path, - extension=source_format, - token=hf_token, - max_files=listing_max, - include_paths=include_paths, - exclude_paths=exclude_paths, - ) - - if is_gasstation: - if target_week: - files = [f for f in files if target_week in f] - logger.info( - f"Filtered to week {target_week} for {dataset_path}: {len(files)} files" - ) - elif num_weeks: - files = filter_files_by_recent_weeks(files, num_weeks) - logger.info( - f"Filtered to last {num_weeks} weeks for {dataset_path}: {len(files)} files" - ) - elif current_week_only: - files = filter_files_by_current_week(files) - logger.info( - f"Filtered to current week files for {dataset_path}: {len(files)} files" - ) - - # Apply max_files limit AFTER week filtering - if needs_week_filter and max_files: - files = files[:max_files] - - return files - - -def _get_download_urls( - dataset_path: str, filenames: List[str], source: str = "huggingface" -) -> List[str]: - """Get download URLs for data files from the specified source. - - Args: - dataset_path: Repository path (org/dataset-name) or S3 path (bucket/prefix) - filenames: List of files to download - source: Source platform ("huggingface", "modelscope", or "s3") - - Returns: - List of download URLs - """ - if source == "modelscope": - return _get_modelscope_urls(dataset_path, filenames) - elif source == "s3": - return _get_s3_urls(dataset_path, filenames) - else: - return _get_huggingface_urls(dataset_path, filenames) - - -def _get_huggingface_urls(dataset_path: str, filenames: List[str]) -> List[str]: - return [ - f"https://huggingface.co/datasets/{dataset_path}/resolve/main/{quote(f, safe='/')}" - for f in filenames - ] - - -def _get_modelscope_urls(dataset_path: str, filenames: List[str]) -> List[str]: - return [ - f"https://www.modelscope.cn/api/v1/datasets/{dataset_path}/repo?Revision=master&FilePath={f}" - for f in filenames - ] - - -def _process_gasstation( - dataset, - to_download: List[str], - temp_dir_root: Path, - media_per_archive: int, - downloaded_archives: Optional[set], - source: str, - hf_token: Optional[str], - seed: Optional[int], -) -> Generator[Dict[str, Any], None, None]: - """Process gasstation datasets. - - Steps: - 1. Download parquet metadata files - 2. Extract unique archive filenames from parquet files - 3. Download those tar archives - 4. Extract media from archives using metadata mapping - - Args: - dataset: BenchmarkDatasetConfig - to_download: List of URLs to download - temp_dir_root: Temporary directory for downloads - media_per_archive: Number of items to extract per archive - downloaded_archives: Set of already-processed parquet basenames - source: Source platform (huggingface, modelscope, s3) - hf_token: HuggingFace API token - seed: Random seed for reproducibility - - Yields: - Sample dictionaries with media and metadata - """ - parquet_urls = [url for url in to_download if ".parquet" in url] - - if not parquet_urls: - logger.warning(f"No parquet files found in download list for {dataset.name}") - return - - # Track parquet files we've already processed - if downloaded_archives is not None: - parquet_basenames = {os.path.basename(url) for url in parquet_urls} - parquets_to_download = [ - url - for url in parquet_urls - if os.path.basename(url) not in downloaded_archives - ] - - skipped = len(parquet_urls) - len(parquets_to_download) - if skipped > 0: - logger.info( - f"Skipping {skipped} already-processed parquet files, " - f"downloading {len(parquets_to_download)} new parquet files" - ) - parquet_urls = parquets_to_download - - if not parquet_urls: - logger.info("No new parquet files to download") - return - - logger.info(f"Downloading {len(parquet_urls)} parquet metadata files") - - max_download_workers = min(10, len(parquet_urls)) - downloaded_parquets = download_files( - parquet_urls, - temp_dir_root, - chunk_size=8192, - max_workers=max_download_workers, - hf_token=hf_token, - ) - - if not downloaded_parquets: - logger.warning(f"Failed to download parquet files for {dataset.name}") - return - - logger.info( - f"Processing {len(downloaded_parquets)} parquet files (per-parquet download & extract)" - ) - - # Build reverse mapping: downloaded filename -> original URL. - # download_single_file() appends a URL hash to filenames (e.g., "file_abc12345.parquet"), - # but downloaded_archives.json must store original basenames (e.g., "file.parquet") - # to match HuggingFace file listings for cache completeness checks. - _downloaded_name_to_url = {} - for url in parquet_urls: - expected_name = _get_expected_download_filename(url) - _downloaded_name_to_url[expected_name] = url - - total_samples = 0 - total_archives = 0 - - # Process each parquet independently: extract metadata ā†’ download its tars ā†’ - # extract samples ā†’ mark parquet done. This ensures a parquet is only marked - # as processed after ALL of its tar archives have been successfully downloaded - # and extracted, and preserves progress incrementally across timeouts. - for parquet_path in downloaded_parquets: - if not parquet_path or not parquet_path.exists(): - continue - - try: - # Resolve original URL basename (not hash-appended) for archive tracking - original_url = _downloaded_name_to_url.get(parquet_path.name) - original_basename = os.path.basename(original_url) if original_url else parquet_path.name - - iso_week = extract_iso_week_from_path(original_url or str(parquet_path)) - - # Step 1: Extract archive filenames referenced by this parquet - archive_filenames = _extract_unique_archive_filenames(parquet_path) - if not archive_filenames: - logger.warning(f"No archive filenames found in {parquet_path}") - # Still mark as done - the parquet was processed, it just had no archives - if downloaded_archives is not None: - downloaded_archives.add(original_basename) - continue - - archive_to_week = {} - for archive_filename in archive_filenames: - if iso_week: - archive_to_week[archive_filename] = iso_week - - # Step 2: Build metadata map from this parquet - metadata_map = _build_parquet_metadata_map(parquet_path) - if iso_week: - for key in metadata_map.keys(): - if key in metadata_map: - metadata_map[key]["iso_week"] = iso_week - - # Step 3: Download this parquet's tar archives - archive_paths = [ - f"archives/{week}/{basename}" - for basename, week in archive_to_week.items() - ] - archive_urls = _get_download_urls(dataset.path, archive_paths, source) - - if not archive_urls: - logger.warning(f"No archive URLs resolved for {original_basename}") - continue - - max_workers = min(10, len(archive_urls)) - downloaded_archive_files = download_files( - archive_urls, - temp_dir_root, - chunk_size=8192, - max_workers=max_workers, - hf_token=hf_token, - ) - - if not downloaded_archive_files: - logger.warning( - f"Failed to download tar archives for parquet {original_basename}" - ) - continue - - # Step 4: Extract samples from this parquet's tar archives - parquet_samples = 0 - for archive_path in downloaded_archive_files: - if not archive_path or not archive_path.exists(): - continue - - try: - tar_iso_week = extract_iso_week_from_path(str(archive_path)) - - for sample in _process_tar_with_metadata( - archive_path, - dataset, - metadata_map, - media_per_archive, - tar_iso_week, - seed, - ): - parquet_samples += 1 - total_samples += 1 - yield sample - except (OSError, pyarrow.ArrowInvalid) as e: - logger.warning(f"Failed to process archive {archive_path}: {e}") - continue - finally: - try: - if archive_path.exists(): - archive_path.unlink() - except OSError: - pass - - total_archives += len(downloaded_archive_files) - - # Step 5: Mark parquet as done ONLY after all its tars are processed. - # This is the last step so if anything above fails or times out, - # the parquet will be retried on the next run. - if downloaded_archives is not None: - downloaded_archives.add(original_basename) - logger.info( - f"āœ… Parquet {original_basename}: {parquet_samples} samples " - f"from {len(downloaded_archive_files)} archives" - ) - except (OSError, pyarrow.ArrowInvalid) as e: - logger.warning(f"Failed to process parquet {parquet_path}: {e}") - continue - finally: - # Clean up this parquet file immediately - try: - if parquet_path and parquet_path.exists(): - parquet_path.unlink() - except (OSError, pyarrow.ArrowInvalid): - pass - - logger.info( - f"Extracted {total_samples} samples from {total_archives} archives" - ) - - -def _extract_unique_archive_filenames(parquet_path: Path) -> set: - """Extract unique archive filenames from gasstation parquet metadata file. - - Args: - parquet_path: Path to the parquet metadata file - - Returns: - Set of unique archive filenames referenced in the parquet - """ - try: - table = pq.read_table(parquet_path) - df = table.to_pandas() - - if "archive_filename" not in df.columns: - logger.warning(f"No archive_filename column in {parquet_path}") - return set() - - archive_filenames = set(df["archive_filename"].dropna().unique()) - return archive_filenames - except (OSError, pyarrow.ArrowInvalid) as e: - logger.warning(f"Failed to extract archive filenames from {parquet_path}: {e}") - return set() - - -def _build_parquet_metadata_map( - parquet_path: Path, -) -> Dict[Tuple[str, str], Dict[str, Any]]: - """Build a mapping from (archive_filename, file_path_in_archive) to metadata. - - Args: - parquet_path: Path to the parquet metadata file - - Returns: - Dictionary mapping (archive_filename, file_path_in_archive) tuples to metadata dicts - """ - try: - table = pq.read_table(parquet_path) - df = table.to_pandas() - - if ( - "archive_filename" not in df.columns - or "file_path_in_archive" not in df.columns - ): - logger.warning(f"Missing required columns in {parquet_path}") - return {} - - metadata_map = {} - for _, row in df.iterrows(): - archive_filename = row.get("archive_filename") - file_path = row.get("file_path_in_archive") - - if pd.isna(archive_filename) or pd.isna(file_path): - continue - - key = (str(archive_filename), str(file_path)) - metadata_map[key] = extract_row_metadata(row, "") - - return metadata_map - except (OSError, json.JSONDecodeError, FileNotFoundError) as e: - logger.warning(f"Failed to build metadata map from {parquet_path}: {e}") - return {} - - -def _process_tar_with_metadata( - archive_path: Path, - dataset, - metadata_map: Dict[Tuple[str, str], Dict[str, Any]], - num_items: int, - iso_week: Optional[str], - seed: Optional[int], -) -> Generator[Dict[str, Any], None, None]: - """Extract media from tar archive using parquet metadata mapping. - - Args: - archive_path: Path to tar archive - dataset: BenchmarkDatasetConfig - metadata_map: Mapping from (archive_filename, file_path_in_archive) to metadata - num_items: Number of items to extract (-1 for all) - iso_week: ISO week string - seed: Random seed - - Yields: - Sample dictionaries with media and metadata - """ - archive_basename = archive_path.name - - archive_basename_normalized = re.sub(r"\.tar_[a-f0-9]{8}\.gz$", ".tar.gz", archive_basename) - - try: - with tarfile.open(archive_path, mode="r:*") as archive: - valid_exts = ( - IMAGE_FILE_EXTENSIONS - if dataset.modality == "image" - else (AUDIO_FILE_EXTENSIONS if dataset.modality == "audio" else VIDEO_FILE_EXTENSIONS) - ) - - all_members = [m for m in archive.getmembers() if m.isreg()] - candidates = [ - m - for m in all_members - if any(m.name.lower().endswith(ext) for ext in valid_exts) - and "MACOSX" not in m.name - ] - - if not candidates: - logger.warning(f"No matching media files found in {archive_path}") - return - - if num_items == -1: - selected = candidates - else: - if seed is not None: - rng = random.Random(seed) - selected = rng.sample(candidates, min(num_items, len(candidates))) - else: - selected = random.sample( - candidates, min(num_items, len(candidates)) - ) - - for member in selected: - try: - src = archive.extractfile(member) - if src is None: - continue - - with closing(src): - data_bytes = src.read() - - if dataset.modality == "image": - try: - media_obj = Image.open(BytesIO(data_bytes)) - except (OSError, FileNotFoundError) as e: - logger.warning(f"Failed to open image {member.name}: {e}") - continue - else: - media_obj = data_bytes - - sample = create_sample(dataset, media_obj, archive_path, iso_week) - sample["archive_filename"] = archive_basename_normalized - sample["member_path"] = member.name - - metadata_key = (archive_basename_normalized, member.name) - if metadata_key in metadata_map: - meta = metadata_map[metadata_key] - for k, v in meta.items(): - if k not in sample: - sample[k] = v - - yield sample - - except Exception as e: - logger.warning( - f"Error extracting {member.name} from {archive_path}: {e}" - ) - continue - - except Exception as e: - logger.warning(f"Error opening tar archive {archive_path}: {e}") - return - - -def _process_zip_or_tar( - source_path: Path, - dataset, - num_items: int, - iso_week: Optional[str] = None, - seed: Optional[int] = None, -): - """Extract media from zip/tar archives (non-gasstation datasets).""" - filename = str(source_path.name).lower() - is_zip = _is_zip_file(filename) - try: - cm = ZipFile(source_path) if is_zip else tarfile.open(source_path, mode="r:*") - with cm as archive: - if dataset.modality == "image": - valid_exts = IMAGE_FILE_EXTENSIONS - elif dataset.modality == "audio": - valid_exts = AUDIO_FILE_EXTENSIONS - else: - valid_exts = VIDEO_FILE_EXTENSIONS - - if is_zip: - list_entries = archive.namelist() - get_name = lambda e: e - - def open_entry(e): - return archive.open(e) - - else: - list_entries = [m for m in archive.getmembers() if m.isreg()] - get_name = lambda m: m.name - - def open_entry(m): - return archive.extractfile(m) - - candidates = [ - e - for e in list_entries - if any(get_name(e).lower().endswith(ext) for ext in valid_exts) - and "MACOSX" not in get_name(e) - ] - if not candidates: - logger.warning(f"No matching files found in {source_path}") - return - - if num_items == -1: - selected = candidates - else: - if seed is not None: - rng = random.Random(seed) - selected = rng.sample(candidates, min(num_items, len(candidates))) - else: - selected = random.sample( - candidates, min(num_items, len(candidates)) - ) - - for entry in selected: - try: - src = open_entry(entry) - if src is None: - continue - with closing(src): - data_bytes = src.read() - - if dataset.modality == "image": - try: - media_obj = Image.open(BytesIO(data_bytes)) - except (OSError, FileNotFoundError): - logger.warning( - f"Failed to open image {get_name(entry)} from {source_path}" - ) - continue - else: - media_obj = data_bytes - - sample = create_sample(dataset, media_obj, source_path, iso_week) - # Normalize hash-appended filename back to original basename - # (download_single_file appends _ before extension) - _normalized = re.sub(r"\.tar_[a-f0-9]{8}\.gz$", ".tar.gz", source_path.name) - if _normalized == source_path.name: - _normalized = re.sub(r"_[a-f0-9]{8}(\.\w+)$", r"\1", source_path.name) - sample["archive_filename"] = _normalized - sample["member_path"] = get_name(entry) - - yield sample - except Exception as e: - logger.warning( - f"Error extracting {get_name(entry)} from {source_path}: {e}" - ) - continue - except Exception as e: - logger.warning(f"Error processing archive file {source_path}: {e}") - return - - -def _process_raw(source_path: Path, dataset, iso_week: Optional[str] = None): - filename = str(source_path.name).lower() - try: - data_bytes = source_path.read_bytes() - if dataset.modality == "image" and any( - filename.endswith(ext) for ext in IMAGE_FILE_EXTENSIONS - ): - media_obj = Image.open(BytesIO(data_bytes)) - elif dataset.modality == "audio" and any( - filename.endswith(ext) for ext in AUDIO_FILE_EXTENSIONS - ): - media_obj = data_bytes - elif dataset.modality == "video" and any( - filename.endswith(ext) for ext in VIDEO_FILE_EXTENSIONS - ): - media_obj = data_bytes - else: - logger.warning( - f"Direct file {source_path} does not match modality {dataset.modality}" - ) - return - - yield create_sample(dataset, media_obj, source_path, iso_week) - except Exception as e: - logger.warning(f"Error reading direct file {source_path}: {e}") - return - - -def _clean_to_json_serializable(value: Any) -> Any: - """Convert arbitrary values to JSON-serializable equivalents.""" - try: - if value is None or isinstance(value, (str, int, bool)): - return value - - if isinstance(value, float): - if np.isnan(value) or np.isinf(value): - return None - return value - - if isinstance(value, (np.integer,)): - return int(value) - if isinstance(value, (np.floating,)): - f = float(value) - if np.isnan(f) or np.isinf(f): - return None - return f - if isinstance(value, (np.bool_,)): - return bool(value) - - if isinstance(value, (bytes, bytearray)): - try: - return base64.b64encode(bytes(value)).decode("ascii") - except Exception: - return None - - if isinstance(value, datetime): - return value.isoformat() - - if isinstance(value, (np.ndarray, list, tuple, set)): - return [ - _clean_to_json_serializable(v) for v in (value.tolist() if isinstance(value, np.ndarray) else list(value)) - ] - - if isinstance(value, dict): - return {str(k): _clean_to_json_serializable(v) for k, v in value.items()} - - try: - json.dumps(value) - return value - except (TypeError, ValueError): - return str(value) - except (TypeError, ValueError): - return None - - -def _extract_parquet_row_metadata(row: Any, media_col: str) -> Dict[str, Any]: - """Extract non-media columns from a pandas Series row and clean to JSON-serializable dict.""" - metadata: Dict[str, Any] = {} - try: - for col, val in row.items(): - if str(col) == str(media_col): - continue - cleaned = _clean_to_json_serializable(val) - col_str = str(col) - metadata[col_str] = cleaned - - col_lower = col_str.lower() - if "hotkey" in col_lower and "generator_hotkey" not in metadata: - metadata["generator_hotkey"] = cleaned - if "uid" in col_lower and "generator_uid" not in metadata: - metadata["generator_uid"] = cleaned - except (OSError, json.JSONDecodeError, FileNotFoundError): - pass - return metadata - - -def _download_filtered_sequential( - dataset, - filenames: List[str], - target_total: int, - temp_dir_root: Path, - source: str, - hf_token: Optional[str], - seed: Optional[int], -) -> Generator[Dict[str, Any], None, None]: - """Download parquet shards one at a time, stopping once target_total samples are collected. - - Filtering is handled inside _process_parquet via dataset.filter_column/filter_value, - so all media decoding follows the exact same path as regular parquet downloads. - """ - remote_paths = _get_download_urls(dataset.path, filenames, source) - random.shuffle(remote_paths) - - collected = 0 - for url in remote_paths: - if collected >= target_total: - break - - logger.info( - f"Filtered download [{dataset.name}]: {collected}/{target_total} collected, " - f"fetching next shard" - ) - downloaded = download_files([url], temp_dir_root, hf_token=hf_token) - if not downloaded: - continue - - parquet_path = downloaded[0] - try: - remaining = target_total - collected - for sample in yield_media_from_source( - parquet_path, dataset, remaining, iso_week=None, - hf_token=hf_token, seed=seed, - ): - yield sample - collected += 1 - if collected >= target_total: - break - finally: - try: - parquet_path.unlink() - except (OSError, pyarrow.ArrowInvalid): - pass - - logger.info( - f"Filtered download complete: {collected}/{target_total} samples for {dataset.name}" - ) - - -def _process_parquet( - source_path: Path, - dataset, - num_items: int, - iso_week: Optional[str] = None, - seed: Optional[int] = None -): - """Process parquet files with embedded media bytes (e.g., PICA-100K). - - Supports datasets with single or multiple image columns. - For datasets with multiple columns (e.g., src_img and tgt_img), yields one sample per column. - """ - try: - table = pq.read_table(source_path) - df = table.to_pandas() - - filter_column = getattr(dataset, "filter_column", None) - filter_value = getattr(dataset, "filter_value", None) - if filter_column and filter_value: - if filter_column not in df.columns: - logger.warning(f"filter_column '{filter_column}' not in {source_path.name}; skipping") - return - df = df[df[filter_column] == filter_value] - if df.empty: - logger.info(f"No rows matching {filter_column}=={filter_value} in {source_path.name}") - return - logger.info(f" {source_path.name}: {len(df)} rows after filter ({filter_column}=={filter_value})") - - if num_items == -1: - sample_df = df - else: - sample_df = df.sample(n=min(num_items, len(df)), random_state=seed) - - data_columns = getattr(dataset, 'data_columns', None) - if data_columns: - media_cols = [c for c in data_columns if c in sample_df.columns] - if not media_cols: - logger.warning( - f"Specified data_columns {data_columns} not found in {source_path}" - ) - return - elif dataset.modality == "image": - media_col = ( - next((c for c in sample_df.columns if c.lower() == "image"), None) - or next((c for c in sample_df.columns if "image" in c.lower() and "_id" not in c.lower()), None) - or next((c for c in sample_df.columns if "image" in c.lower()), None) - ) - media_cols = [media_col] if media_col else [] - elif dataset.modality == "audio": - candidates = ["audio", "bytes", "content", "data", "wav", "mp3"] - exact = [c for c in sample_df.columns if c.lower() == "audio"] - if exact: - media_cols = exact - else: - media_cols = [ - c - for c in sample_df.columns - if "_id" not in c.lower() - and any(k in c.lower() for k in candidates) - ] - else: - candidates = ["video", "bytes", "content", "data"] - media_col = ( - next((c for c in sample_df.columns if c.lower() in candidates), None) - or next((c for c in sample_df.columns if any(k in c.lower() for k in candidates) and "_id" not in c.lower()), None) - or next((c for c in sample_df.columns if any(k in c.lower() for k in candidates)), None) - ) - media_cols = [media_col] if media_col else [] - - if not media_cols: - logger.warning( - f"No media column found in {source_path} for modality {dataset.modality}" - ) - return - - for _, row in sample_df.iterrows(): - for col in media_cols: - try: - media_data = row[col] - audio_sampling_rate = None # Track sampling rate for audio dicts - - if isinstance(media_data, dict): - # For audio, try to extract sampling_rate before getting the array - if dataset.modality == "audio": - audio_sampling_rate = media_data.get("sampling_rate") or media_data.get("sample_rate") - - key = next( - ( - k - for k in media_data - if any( - s in k.lower() - for s in ["bytes", "image", "video", "audio", "array", "data", "content"] - ) - ), - None, - ) - if key: - media_data = media_data[key] - else: - logger.warning(f"No valid key found in dict media_data for {source_path}: {list(media_data.keys())}") - continue - - if dataset.modality == "image": - if media_data is None or isinstance(media_data, (int, float)): - continue - - try: - img = Image.open(BytesIO(media_data)) - except (OSError, FileNotFoundError): - if isinstance(media_data, str): - media_data = base64.b64decode(media_data) - img = Image.open(BytesIO(media_data)) - sample = create_sample(dataset, img, source_path, iso_week) - else: - if media_data is None or isinstance(media_data, (int, float)): - continue - - # Handle audio array format (common in HF audio datasets) - if dataset.modality == "audio" and isinstance(media_data, (list, np.ndarray)): - # Audio data is a numpy array - convert to WAV bytes - try: - audio_array = np.array(media_data, dtype=np.float32) - if audio_array.ndim == 1: - audio_array = audio_array.reshape(-1, 1) - sr = audio_sampling_rate if audio_sampling_rate else 16000 - buffer = io.BytesIO() - sf.write(buffer, audio_array, sr, format='WAV') - media_data = buffer.getvalue() - except (ValueError, TypeError, OSError) as e: - logger.warning(f"Failed to convert audio array to bytes: {e}") - continue - elif not isinstance(media_data, (bytes, bytearray)): - if isinstance(media_data, str): - media_data = base64.b64decode(media_data) - else: - continue - sample = create_sample(dataset, bytes(media_data), source_path, iso_week) - - row_metadata = _extract_parquet_row_metadata(row, col) - for k, v in row_metadata.items(): - if k not in sample: - sample[k] = v - - if len(media_cols) > 1: - sample["source_column"] = col - - yield sample - except Exception as e: - logger.warning(f"Failed to extract row from {source_path} (column {col}): {e}") - continue - except Exception as e: - logger.warning(f"Error processing parquet file {source_path}: {e}") - return - - -def _process_frame_directory( - frame_dir: Path, dataset, iso_week: Optional[str] = None -) -> Generator[Dict[str, Any], None, None]: - """Process a directory containing pre-extracted video frames. - - Args: - frame_dir: Directory containing frame files (e.g., bm-videos/dfb/DFDC/test/frames/amwhgrjvkw/) - dataset: BenchmarkDatasetConfig - iso_week: Optional ISO week string - - Yields: - Sample dict with 'video_frames' key containing list of frame paths - """ - try: - if not frame_dir.is_dir(): - logger.warning(f"{frame_dir} is not a directory") - return - - # Get all image files in the directory - frame_files = [] - for ext in IMAGE_FILE_EXTENSIONS: - frame_files.extend(sorted(frame_dir.glob(f"*{ext}"))) - frame_files.extend(sorted(frame_dir.glob(f"*{ext.upper()}"))) - - # Sort frames numerically if they have numeric names - def extract_number(filepath): - """Extract numeric part from filename for sorting (e.g., '000.png' -> 0)""" - name = filepath.stem - match = re.match(r"(\d+)", name) - if match: - return int(match.group(1)) - return name - - frame_files = sorted(set(frame_files), key=extract_number) - - if not frame_files: - logger.warning(f"No frame files found in {frame_dir}") - return - - logger.debug(f"Found {len(frame_files)} frames in {frame_dir.name}") - - # Create sample with frame paths - sample = { - "video_frames": frame_files, - "media_type": dataset.media_type, - "dataset_name": dataset.name, - "dataset_path": dataset.path, - "source_file": frame_dir.name, - } - - if iso_week: - sample["iso_week"] = iso_week - - yield sample - - except Exception as e: - logger.warning(f"Error processing frame directory {frame_dir}: {e}") - return - - -def list_hf_files( - repo_id, - repo_type="dataset", - extension=None, - token=None, - max_files=None, - include_paths=None, - exclude_paths=None, -): - """List files from a Hugging Face repository with early termination support. - - Args: - repo_id: Repository ID - repo_type: Type of repository ('dataset', 'model', etc.) - extension: Filter files by extension - token: Hugging Face API token for private datasets - max_files: Stop after collecting this many matching files (None = no limit) - include_paths: Only include files containing one of these path segments - exclude_paths: Exclude files containing any of these path segments - - Returns: - List of files in the repository - - Raises: - DatasetAccessError: If the repository is gated/private without access, or not found - """ - files = [] - if extension: - if isinstance(extension, (list, tuple, set)): - exts = tuple(extension) - else: - exts = (extension,) - else: - exts = None - - try: - for f in hf_hub.list_repo_files(repo_id=repo_id, repo_type=repo_type, token=token): - if exts and not f.endswith(exts): - continue - if include_paths and not any(path_seg in f for path_seg in include_paths): - continue - if exclude_paths and any(path_seg in f for path_seg in exclude_paths): - continue - files.append(f) - if max_files and len(files) >= max_files: - logger.info(f"Early termination: collected {max_files} files from {repo_id}") - break - except GatedRepoError: - raise DatasetAccessError( - f"Dataset {repo_id} is gated and requires access approval. " - "Visit the dataset page on HuggingFace to request access." - ) - except RepositoryNotFoundError: - raise DatasetAccessError( - f"Dataset {repo_id} not found. It may not exist or may be private." - ) - except Exception as e: - logger.error(f"Failed to list files of type {extension} in {repo_id}: {e}") - return files - - -def list_modelscope_files(repo_id, extension=None): - """List files from a ModelScope dataset repository. - - Args: - repo_id: Repository ID (format: org/dataset-name) - extension: Filter files by extension(s) - - Returns: - List of files in the repository - """ - files = [] - try: - api = MSHubApi() - file_info = api.get_dataset_files(repo_id, revision="master", recursive=True) - - # Extract file paths - if isinstance(file_info, list): - files = [ - f["Path"] if isinstance(f, dict) and "Path" in f else str(f) - for f in file_info - ] - elif isinstance(file_info, dict): - files = [f for f in file_info.keys()] if file_info else [] - - if extension and files: - if isinstance(extension, (list, tuple, set)): - exts = tuple(extension) - files = [f for f in files if f.endswith(exts)] - else: - files = [f for f in files if f.endswith(extension)] - - except Exception as e: - logger.error( - f"Failed to list files of type {extension} in ModelScope repo {repo_id}: {e}" - ) - logger.error( - f"Make sure 'modelscope' package is installed: pip install modelscope" - ) - - return files - - -def _is_zip_file(filename_lower: str) -> bool: - """Return True if filename looks like a zip archive.""" - return filename_lower.endswith(".zip") - - -def _is_tar_file(filename_lower: str) -> bool: - """Return True if filename looks like a tar archive (.tar, .tar.gz, .tgz). - - Also handles hash-suffixed filenames like 'file.tar_abc123.gz' that result - from download_single_file adding URL hashes to avoid collisions. - """ - if filename_lower.endswith(".tar") or filename_lower.endswith(".tgz"): - return True - if filename_lower.endswith(".tar.gz"): - return True - # Handle hash-suffixed tar.gz files: file.tar_.gz - if re.search(r"\.tar_[a-f0-9]{8}\.gz$", filename_lower): - return True - return False - - -def _is_parquet_file(filename_lower: str) -> bool: - """Return True if filename looks like a parquet file.""" - return filename_lower.endswith(".parquet") - - -def _stream_downloads( - urls: List[str], - output_dir: Path, - chunk_size: int = 8192, - max_workers: int = 10, - hf_token: Optional[str] = None, -) -> Generator[Optional[Path], None, None]: - """Like download_files but yields each path as soon as its download completes.""" - from concurrent.futures import ThreadPoolExecutor, as_completed - - output_dir.mkdir(parents=True, exist_ok=True) - if not urls: - return - - def download_url(url): - try: - return download_single_file(url, output_dir, chunk_size, hf_token) - except (requests.RequestException, OSError) as e: - logger.error(f"Error downloading {url}: {e}") - return None - - workers = min(max_workers, len(urls)) - with ThreadPoolExecutor(max_workers=workers) as executor: - futures = {executor.submit(download_url, url): url for url in urls} - for future in as_completed(futures): - yield future.result() - - -def download_files( - urls: List[str], - output_dir: Path, - chunk_size: int = 8192, - max_workers: int = 10, - hf_token: Optional[str] = None, -) -> List[Path]: - """Download multiple files in parallel. - - Args: - urls: List of URLs to download - output_dir: Directory to save the files - chunk_size: Size of chunks to download at a time - max_workers: Maximum number of parallel downloads (default: 10) - hf_token: Hugging Face API token for private datasets - - Returns: - List of successfully downloaded file paths - """ - from concurrent.futures import ThreadPoolExecutor, as_completed - - output_dir.mkdir(parents=True, exist_ok=True) - downloaded_files = [] - - def download_url(url): - try: - return download_single_file(url, output_dir, chunk_size, hf_token) - except (requests.RequestException, OSError) as e: - logger.error(f"Error downloading {url}: {e}") - return None - - max_workers = min(max_workers, len(urls)) if urls else 1 - - logger.debug(f"Downloading {len(urls)} files with {max_workers} parallel workers") - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(download_url, url): url for url in urls} - - for future in as_completed(futures): - result = future.result() - if result: - downloaded_files.append(result) - - return downloaded_files - - -def download_single_file( - url: str, output_dir: Path, chunk_size: int, hf_token: Optional[str] = None -) -> Optional[Path]: - """Download a single file synchronously with resume and retry support - - Args: - url: URL to download (supports http/https and s3: schemes) - output_dir: Directory to save the file - chunk_size: Size of chunks to download at a time - hf_token: Hugging Face API token for private datasets - - Returns: - Path to the downloaded file, or None if failed - """ - try: - # Handle S3 URLs (format: s3:bucket-name/key/path or s3:bucket-name/frame/directory/) - if url.startswith("s3:"): - from .utils.s3_utils import download_s3_frame_directory - - s3_path = url[3:] # Remove "s3:" prefix - parts = s3_path.split("/", 1) - if len(parts) != 2: - logger.error( - f"Invalid S3 URL format: {url}. Expected: s3:bucket/key/path" - ) - return None - - bucket = parts[0] - key = parts[1] - - # Check if this is a frame directory (no file extension or ends with /) - if key.endswith("/") or ("." not in os.path.basename(key)): - # This is a frame directory - download all frames - return download_s3_frame_directory(bucket, key, output_dir) - else: - # This is a single file - use unique filename to avoid collisions - url_hash = hashlib.md5(key.encode()).hexdigest()[:8] - base_filename = os.path.basename(key) - name, ext = os.path.splitext(base_filename) - filename = f"{name}_{url_hash}{ext}" - filepath = output_dir / filename - return download_s3_file(bucket, key, filepath) - - # Handle regular HTTP/HTTPS URLs - # Use hash of full URL path to avoid filename collisions when multiple - # files from different directories have the same basename - url_hash = hashlib.md5(url.encode()).hexdigest()[:8] - base_filename = os.path.basename(url) - name, ext = os.path.splitext(base_filename) - filename = f"{name}_{url_hash}{ext}" - filepath = output_dir / filename - partial_filepath = Path(str(filepath) + ".partial") - - logger.debug(f"Downloading {url}") - - max_retries = 5 - effective_chunk_size = max(chunk_size, 1024 * 1024) - - for attempt in range(max_retries): - try: - headers = {} - if hf_token and "huggingface.co" in url: - headers["Authorization"] = f"Bearer {hf_token}" - - downloaded_size = 0 - if partial_filepath.exists(): - downloaded_size = partial_filepath.stat().st_size - headers["Range"] = f"bytes={downloaded_size}-" - logger.info(f"Resuming download from {downloaded_size} bytes") - - response = requests.get( - url, stream=True, timeout=(30, 300), headers=headers - ) - - if response.status_code == 416: - logger.info("Server returned 416 (range not satisfiable), restarting download") - partial_filepath.unlink(missing_ok=True) - downloaded_size = 0 - del headers["Range"] - response = requests.get( - url, stream=True, timeout=(30, 300), headers=headers - ) - - if response.status_code not in (200, 206): - logger.error(f"Failed to download {url}: Status {response.status_code}") - return None - - if response.status_code == 200 and downloaded_size > 0: - logger.info("Server doesn't support resume, restarting download") - partial_filepath.unlink(missing_ok=True) - downloaded_size = 0 - - total_size = int(response.headers.get("content-length", 0)) - if response.status_code == 200: - expected_total = total_size - else: - expected_total = downloaded_size + total_size - - mode = "ab" if downloaded_size > 0 else "wb" - with open(partial_filepath, mode) as f: - for chunk in response.iter_content(chunk_size=effective_chunk_size): - if chunk: - f.write(chunk) - - actual_size = partial_filepath.stat().st_size - if expected_total > 0 and actual_size != expected_total: - raise IOError( - f"Download incomplete: expected {expected_total} bytes, got {actual_size} bytes" - ) - - partial_filepath.rename(filepath) - return filepath - - except (requests.exceptions.RequestException, IOError) as e: - wait_time = min(30 * (2 ** attempt), 300) - logger.warning( - f"Download attempt {attempt + 1}/{max_retries} failed: {e}. " - f"Retrying in {wait_time}s..." - ) - if attempt < max_retries - 1: - time.sleep(wait_time) - else: - raise - - return None - - except Exception as e: - logger.error(f"Error downloading {url}: {str(e)}") - logger.error(traceback.format_exc()) - return None - - -def _get_expected_download_filename(url: str) -> str: - """Compute the filename that download_single_file() would produce for a URL. - - download_single_file() appends a URL hash to filenames to avoid collisions - (e.g., "file.parquet" becomes "file_abc12345.parquet"). This function mirrors - that naming convention to enable reverse-mapping from downloaded filenames - back to original URLs. - """ - if url.startswith("s3:"): - s3_path = url[3:] - parts = s3_path.split("/", 1) - key = parts[1] if len(parts) == 2 else s3_path - url_hash = hashlib.md5(key.encode()).hexdigest()[:8] - base_filename = os.path.basename(key) - else: - url_hash = hashlib.md5(url.encode()).hexdigest()[:8] - base_filename = os.path.basename(url) - name, ext = os.path.splitext(base_filename) - return f"{name}_{url_hash}{ext}" - - -def _is_dataset_cached(dataset, cache_dir: str = "/.cache/gasbench") -> bool: - """ - Check if a specific dataset is already cached locally. - - Returns True if the dataset appears to be cached and has samples available. - """ - try: - dataset_dir = f"{cache_dir}/datasets/{dataset.name}" - - dataset_info_file = os.path.join(dataset_dir, "dataset_info.json") - samples_dir = os.path.join(dataset_dir, "samples") - metadata_file = os.path.join(dataset_dir, "sample_metadata.json") - - # Check if all required files exist and samples directory has content - return ( - os.path.exists(dataset_info_file) - and os.path.exists(samples_dir) - and os.path.exists(metadata_file) - and len(os.listdir(samples_dir)) > 0 - ) - except (OSError, json.JSONDecodeError, FileNotFoundError) as e: - logger.warning(f"Error checking cache for dataset {dataset.name}: {e}") - return False - - -def _load_dataset_from_cache(dataset, cache_dir: str = "/.cache/gasbench"): - """ - Load samples from a cached dataset locally. - - Yields samples in the same format as download_and_extract. - """ - - dataset_dir = f"{cache_dir}/datasets/{dataset.name}" - samples_dir = os.path.join(dataset_dir, "samples") - metadata_file = os.path.join(dataset_dir, "sample_metadata.json") - - # Load sample metadata - with open(metadata_file, "r") as f: - sample_metadata = json.load(f) - - sample_files = [f for f in os.listdir(samples_dir) if not f.startswith(".")] - - logger.info(f"Loading {len(sample_files)} cached samples for {dataset.name}") - - for filename in sample_files: - try: - metadata = sample_metadata.get(filename, {}) - filepath = os.path.join(samples_dir, filename) - - if dataset.modality == "image": - # Load image file - img = Image.open(filepath) - sample = { - "image": img, - "media_type": dataset.media_type, - "dataset_name": dataset.name, - "dataset_path": dataset.path, - "source_file": f"cached_{filename}", - **metadata, - } - yield sample - - elif dataset.modality == "audio": - # Load audio file as bytes - with open(filepath, "rb") as f: - audio_bytes = f.read() - sample = { - "audio_bytes": audio_bytes, - "media_type": dataset.media_type, - "dataset_name": dataset.name, - "dataset_path": dataset.path, - "source_file": f"cached_{filename}", - **metadata, # Include all original metadata - } - yield sample - - elif dataset.modality == "video": - # Check if it's a frame directory or a video file - if os.path.isdir(filepath): - # Frame directory - load all frames - frame_files = [] - for ext in IMAGE_FILE_EXTENSIONS: - frame_files.extend(sorted(Path(filepath).glob(f"*{ext}"))) - frame_files.extend( - sorted(Path(filepath).glob(f"*{ext.upper()}")) - ) - - frame_files = sorted(set(frame_files), key=lambda p: p.name) - - sample = { - "video_frames": frame_files, - "media_type": dataset.media_type, - "dataset_name": dataset.name, - "dataset_path": dataset.path, - "source_file": f"cached_{filename}", - **metadata, - } - yield sample - else: - # Regular video file - with open(filepath, "rb") as f: - video_bytes = f.read() - - sample = { - "video_bytes": video_bytes, - "media_type": dataset.media_type, - "dataset_name": dataset.name, - "dataset_path": dataset.path, - "source_file": f"cached_{filename}", - **metadata, - } - yield sample - - except (OSError, json.JSONDecodeError, FileNotFoundError) as e: - logger.warning( - f"Failed to load cached sample {filename} from {dataset.name}: {e}" - ) - continue diff --git a/src/gasbench/dataset/download/__init__.py b/src/gasbench/dataset/download/__init__.py new file mode 100644 index 0000000..9825390 --- /dev/null +++ b/src/gasbench/dataset/download/__init__.py @@ -0,0 +1,39 @@ +"""Dataset download and extraction. + +Previously a single ~1900-line module; split into a package with the same +public API (all prior `from ...dataset.download import X` imports still work): + + constants - shared constants, DatasetAccessError, archive-type predicates + fetch - low-level file downloading (HTTP/HF streaming) + listing - remote file enumeration and download-URL resolution + extract - decode media samples from parquet/tar/zip/raw sources + cache_io - load samples from the persistent cache volume + core - download_and_extract orchestration (incl. gasstation) +""" + +from .constants import ( + DatasetAccessError, + IMAGE_FILE_EXTENSIONS, + VIDEO_FILE_EXTENSIONS, + AUDIO_FILE_EXTENSIONS, + MAX_FILES_DEFAULT, +) +from .fetch import download_files, download_single_file +from .listing import list_hf_files, list_modelscope_files +from .extract import yield_media_from_source, _process_parquet +from .core import download_and_extract + +__all__ = [ + "download_and_extract", + "DatasetAccessError", + "IMAGE_FILE_EXTENSIONS", + "VIDEO_FILE_EXTENSIONS", + "AUDIO_FILE_EXTENSIONS", + "MAX_FILES_DEFAULT", + "download_files", + "download_single_file", + "list_hf_files", + "list_modelscope_files", + "yield_media_from_source", + "_process_parquet", +] diff --git a/src/gasbench/dataset/download/cache_io.py b/src/gasbench/dataset/download/cache_io.py new file mode 100644 index 0000000..148d215 --- /dev/null +++ b/src/gasbench/dataset/download/cache_io.py @@ -0,0 +1,135 @@ +"""Load previously-cached dataset samples from the persistent cache volume.""" + +import json +import os +from pathlib import Path + +from PIL import Image + +from ...logger import get_logger + +from .constants import IMAGE_FILE_EXTENSIONS + +logger = get_logger(__name__) + + +def _is_dataset_cached(dataset, cache_dir: str = "/.cache/gasbench") -> bool: + """ + Check if a specific dataset is already cached locally. + + Returns True if the dataset appears to be cached and has samples available. + """ + try: + dataset_dir = f"{cache_dir}/datasets/{dataset.name}" + + dataset_info_file = os.path.join(dataset_dir, "dataset_info.json") + samples_dir = os.path.join(dataset_dir, "samples") + metadata_file = os.path.join(dataset_dir, "sample_metadata.json") + + # Check if all required files exist and samples directory has content + return ( + os.path.exists(dataset_info_file) + and os.path.exists(samples_dir) + and os.path.exists(metadata_file) + and len(os.listdir(samples_dir)) > 0 + ) + except (OSError, json.JSONDecodeError, FileNotFoundError) as e: + logger.warning(f"Error checking cache for dataset {dataset.name}: {e}") + return False + + + +def _load_dataset_from_cache(dataset, cache_dir: str = "/.cache/gasbench"): + """ + Load samples from a cached dataset locally. + + Yields samples in the same format as download_and_extract. + """ + + dataset_dir = f"{cache_dir}/datasets/{dataset.name}" + samples_dir = os.path.join(dataset_dir, "samples") + metadata_file = os.path.join(dataset_dir, "sample_metadata.json") + + # Load sample metadata + with open(metadata_file, "r") as f: + sample_metadata = json.load(f) + + sample_files = [f for f in os.listdir(samples_dir) if not f.startswith(".")] + + logger.info(f"Loading {len(sample_files)} cached samples for {dataset.name}") + + for filename in sample_files: + try: + metadata = sample_metadata.get(filename, {}) + filepath = os.path.join(samples_dir, filename) + + if dataset.modality == "image": + # Load image file + img = Image.open(filepath) + sample = { + "image": img, + "media_type": dataset.media_type, + "dataset_name": dataset.name, + "dataset_path": dataset.path, + "source_file": f"cached_{filename}", + **metadata, + } + yield sample + + elif dataset.modality == "audio": + # Load audio file as bytes + with open(filepath, "rb") as f: + audio_bytes = f.read() + sample = { + "audio_bytes": audio_bytes, + "media_type": dataset.media_type, + "dataset_name": dataset.name, + "dataset_path": dataset.path, + "source_file": f"cached_{filename}", + **metadata, # Include all original metadata + } + yield sample + + elif dataset.modality == "video": + # Check if it's a frame directory or a video file + if os.path.isdir(filepath): + # Frame directory - load all frames + frame_files = [] + for ext in IMAGE_FILE_EXTENSIONS: + frame_files.extend(sorted(Path(filepath).glob(f"*{ext}"))) + frame_files.extend( + sorted(Path(filepath).glob(f"*{ext.upper()}")) + ) + + frame_files = sorted(set(frame_files), key=lambda p: p.name) + + sample = { + "video_frames": frame_files, + "media_type": dataset.media_type, + "dataset_name": dataset.name, + "dataset_path": dataset.path, + "source_file": f"cached_{filename}", + **metadata, + } + yield sample + else: + # Regular video file + with open(filepath, "rb") as f: + video_bytes = f.read() + + sample = { + "video_bytes": video_bytes, + "media_type": dataset.media_type, + "dataset_name": dataset.name, + "dataset_path": dataset.path, + "source_file": f"cached_{filename}", + **metadata, + } + yield sample + + except (OSError, json.JSONDecodeError, FileNotFoundError) as e: + logger.warning( + f"Failed to load cached sample {filename} from {dataset.name}: {e}" + ) + continue + diff --git a/src/gasbench/dataset/download/constants.py b/src/gasbench/dataset/download/constants.py new file mode 100644 index 0000000..e3f73d5 --- /dev/null +++ b/src/gasbench/dataset/download/constants.py @@ -0,0 +1,43 @@ +"""Shared constants, the dataset-access error, and archive-type predicates.""" + +import re + +class DatasetAccessError(Exception): + """Raised when a dataset cannot be accessed (gated, not found, or permission denied).""" + + pass + + +IMAGE_FILE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"} +VIDEO_FILE_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".webm", ".m4v", ".mpeg", ".mpg"} +AUDIO_FILE_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac"} + +MAX_FILES_DEFAULT = 5000 + +def _is_zip_file(filename_lower: str) -> bool: + """Return True if filename looks like a zip archive.""" + return filename_lower.endswith(".zip") + + +def _is_tar_file(filename_lower: str) -> bool: + """Return True if filename looks like a tar archive (.tar, .tar.gz, .tgz). + + Also handles hash-suffixed filenames like 'file.tar_abc123.gz' that result + from download_single_file adding URL hashes to avoid collisions. + """ + if filename_lower.endswith(".tar") or filename_lower.endswith(".tgz"): + return True + if filename_lower.endswith(".tar.gz"): + return True + # Handle hash-suffixed tar.gz files: file.tar_.gz + if re.search(r"\.tar_[a-f0-9]{8}\.gz$", filename_lower): + return True + return False + + +def _is_parquet_file(filename_lower: str) -> bool: + """Return True if filename looks like a parquet file.""" + return filename_lower.endswith(".parquet") + + + diff --git a/src/gasbench/dataset/download/core.py b/src/gasbench/dataset/download/core.py new file mode 100644 index 0000000..089a0e5 --- /dev/null +++ b/src/gasbench/dataset/download/core.py @@ -0,0 +1,617 @@ +"""Top-level download orchestration: download_and_extract, gasstation, filtered-sequential.""" + +import json +import os +import random +import shutil +import tempfile +from pathlib import Path +from typing import Any, Dict, Generator, List, Optional + +import pyarrow + +from ...logger import get_logger +from ..utils.gasstation_utils import ( + extract_iso_week_from_path, +) + +from .constants import ( + DatasetAccessError, + IMAGE_FILE_EXTENSIONS, + VIDEO_FILE_EXTENSIONS, + AUDIO_FILE_EXTENSIONS, +) +from .cache_io import _is_dataset_cached, _load_dataset_from_cache +from .listing import ( + _list_remote_dataset_files, + _get_download_urls, + _select_files_to_download, +) +from .fetch import _stream_downloads, download_files, _get_expected_download_filename +from .extract import ( + yield_media_from_source, + _extract_unique_archive_filenames, + _build_parquet_metadata_map, + _process_tar_with_metadata, +) + +logger = get_logger(__name__) + + +def _calculate_files_to_download( + dataset, + source_format: str, + media_per_archive: int, + archives_per_dataset: int, +) -> int: + """Calculate # files to download based on dataset modality and source format. + Returns -1 to indicate "download all files", or a positive integer for the count. + """ + src_fmt = source_format.lower().lstrip(".") + + # Direct media files (jpg, png, mp4, etc.) - download equivalent to archive extraction + if dataset.modality == "image": + is_direct_media = src_fmt in {ext.lstrip(".") for ext in IMAGE_FILE_EXTENSIONS} + elif dataset.modality == "audio": + is_direct_media = src_fmt in {ext.lstrip(".") for ext in AUDIO_FILE_EXTENSIONS} + else: + is_direct_media = src_fmt in {ext.lstrip(".") for ext in VIDEO_FILE_EXTENSIONS} + + if is_direct_media: + if media_per_archive == -1 or archives_per_dataset == -1: + return -1 + return media_per_archive * archives_per_dataset + + # Archive-based media (parquet, zip, tar, etc.) + return archives_per_dataset + + + +def download_and_extract( + dataset, # BenchmarkDatasetConfig + media_per_archive: int = 100, + archives_per_dataset: int = 5, + temp_dir: Optional[str] = None, + force_download: bool = False, + current_week_only: bool = False, + num_weeks: int = None, + cache_dir: str = "/.cache/gasbench", + downloaded_archives: Optional[set] = None, + target_week: Optional[str] = None, + hf_token: Optional[str] = None, + seed: Optional[int] = None, +) -> Generator[Dict[str, Any], None, None]: + """ + Download datasets and yield extracted media as a generator. + + Downloads files temporarily, extracts/processing content, then cleans up downloads. + Processed content (videos, images) is cached persistently in other locations. + + Args: + dataset: BenchmarkDatasetConfig object + media_per_archive: Number of media items (images/videos) to extract per archive file + archives_per_dataset: Number of archive files to download per dataset + temp_dir: Temporary directory for downloads + force_download: Force download even if dataset appears to be cached + cache_dir: Base cache directory for persistent storage + + Yields: + Complete sample dictionaries ready for processing functions + """ + try: + if not force_download and _is_dataset_cached(dataset, cache_dir): + logger.info(f"Dataset {dataset.name} found in cache, loading from volume") + try: + cached_samples_count = 0 + for sample in _load_dataset_from_cache(dataset, cache_dir): + cached_samples_count += 1 + yield sample + logger.info( + f"Successfully loaded {cached_samples_count} samples from cache for {dataset.name}" + ) + return + except (OSError, json.JSONDecodeError, FileNotFoundError) as e: + logger.warning(f"Failed to load {dataset.name} from cache: {e}") + logger.info( + f"Fallback to download: {dataset.name} will be downloaded fresh" + ) + # Fall through to download logic + + # Create temporary directory for downloads + if temp_dir is not None: + try: + Path(temp_dir).mkdir(parents=True, exist_ok=True) + except OSError: + pass + temp_dir_root = Path(tempfile.mkdtemp(dir=temp_dir)) + + try: + include_paths = getattr(dataset, "include_paths", None) + exclude_paths = getattr(dataset, "exclude_paths", None) + source = getattr(dataset, "source", "huggingface") + + n_files = _calculate_files_to_download( + dataset, + dataset.source_format, + media_per_archive, + archives_per_dataset, + ) + max_files_to_list = n_files if n_files > 0 else None + + # ā”€ā”€ Filtered sequential mode ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + # When filter_column/filter_value are set, download parquet shards + # one at a time and stop as soon as the sample target is hit. + filter_column = getattr(dataset, "filter_column", None) + filter_value = getattr(dataset, "filter_value", None) + if filter_column and filter_value: + try: + all_filenames = _list_remote_dataset_files( + dataset.path, + dataset.source_format or "parquet", + False, None, None, + include_paths, exclude_paths, + source, hf_token, + max_files=None, # need full list to shuffle and pick from + ) + except DatasetAccessError as e: + logger.warning(f"Skipping {dataset.name}: {e}") + return + + if not all_filenames: + logger.warning(f"No parquet files found for {dataset.path}") + return + + target_total = media_per_archive * archives_per_dataset + logger.info( + f"Filtered mode: {dataset.name} ā€” {filter_column}=={filter_value}, " + f"target {target_total} samples from {len(all_filenames)} shards" + ) + yield from _download_filtered_sequential( + dataset=dataset, + filenames=all_filenames, + target_total=target_total, + temp_dir_root=temp_dir_root, + source=source, + hf_token=hf_token, + seed=seed, + ) + return + # ā”€ā”€ Standard (non-filtered) download path ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ā”€ + + try: + filenames = _list_remote_dataset_files( + dataset.path, + dataset.source_format, + current_week_only, + num_weeks, + target_week, + include_paths, + exclude_paths, + source, + hf_token, + max_files=max_files_to_list, + ) + except DatasetAccessError as e: + logger.warning( + f"Skipping dataset {dataset.name}: {e}" + ) + return + + if not filenames: + is_gasstation = "gasstation" in dataset.name.lower() + if is_gasstation: + logger.warning( + f"No files found for {dataset.path} with format {dataset.source_format}. " + f"Gasstation datasets require parquet metadata files." + ) + return + + logger.warning( + f"No files found for {dataset.path} with format {dataset.source_format}" + ) + + fallback_formats = [".parquet", ".zip", ".tar", ".tar.gz"] + for fallback_format in fallback_formats: + if fallback_format != dataset.source_format: + logger.info( + f"Trying fallback format {fallback_format} for {dataset.path}" + ) + try: + filenames = _list_remote_dataset_files( + dataset.path, + fallback_format, + current_week_only, + num_weeks, + target_week, + include_paths, + exclude_paths, + source, + hf_token, + max_files=max_files_to_list, + ) + except DatasetAccessError as e: + logger.warning( + f"Skipping dataset {dataset.name}: {e}" + ) + return + if filenames: + logger.info( + f"Found {len(filenames)} files with format {fallback_format}" + ) + # Recalculate n_files for the actual format found ā€” the + # original n_files was computed for source_format (e.g. mp4) + # which treats each file as one media item, but an archive + # format (parquet/zip/tar) should only download + # archives_per_dataset files, not media_per_archive. + n_files = _calculate_files_to_download( + dataset, fallback_format, media_per_archive, archives_per_dataset + ) + break + + if not filenames: + logger.warning(f"No files found for {dataset.path} with any format") + return + + remote_paths = _get_download_urls(dataset.path, filenames, source) + + is_gasstation = "gasstation" in dataset.name.lower() + to_download = _select_files_to_download( + remote_paths, n_files, prioritize_recent=is_gasstation, seed=seed + ) + + if is_gasstation and downloaded_archives is not None: + original_count = len(to_download) + to_download = [ + url + for url in to_download + if os.path.basename(url) not in downloaded_archives + ] + filtered_count = original_count - len(to_download) + if filtered_count > 0: + logger.info( + f"Skipping {filtered_count} already-downloaded archives, " + f"downloading {len(to_download)} new archives" + ) + + if len(to_download) == 0: + return + + logger.info( + f"Downloading {len(to_download)} files from {dataset.path} (dataset: {dataset.name})" + ) + + if ( + is_gasstation + and to_download + and any(".parquet" in url for url in to_download) + ): + yield from _process_gasstation( + dataset=dataset, + to_download=to_download, + temp_dir_root=temp_dir_root, + media_per_archive=media_per_archive, + downloaded_archives=downloaded_archives, + source=source, + hf_token=hf_token, + seed=seed, + ) + else: + max_download_workers = min(10, len(to_download)) + + successfully_processed = 0 + any_downloaded = False + for downloaded_file in _stream_downloads( + to_download, + temp_dir_root, + chunk_size=8192, + max_workers=max_download_workers, + hf_token=hf_token, + ): + if not downloaded_file or not downloaded_file.exists(): + continue + + any_downloaded = True + try: + iso_week = extract_iso_week_from_path(str(downloaded_file)) + successfully_processed += 1 + + for sample in yield_media_from_source( + downloaded_file, + dataset, + media_per_archive, + iso_week, + hf_token, + seed, + ): + yield sample + finally: + try: + if downloaded_file.exists(): + downloaded_file.unlink() + except OSError: + pass + + if not any_downloaded: + logger.warning( + f"No files successfully downloaded for {dataset.name}" + ) + return + + logger.info( + f"Downloaded and processed {successfully_processed} files " + f"for {dataset.name}" + ) + + finally: + if temp_dir_root.exists(): + shutil.rmtree(temp_dir_root) + + except Exception as e: + logger.error(f"Error processing {dataset.path}: {e}") + + + +def _process_gasstation( + dataset, + to_download: List[str], + temp_dir_root: Path, + media_per_archive: int, + downloaded_archives: Optional[set], + source: str, + hf_token: Optional[str], + seed: Optional[int], +) -> Generator[Dict[str, Any], None, None]: + """Process gasstation datasets. + + Steps: + 1. Download parquet metadata files + 2. Extract unique archive filenames from parquet files + 3. Download those tar archives + 4. Extract media from archives using metadata mapping + + Args: + dataset: BenchmarkDatasetConfig + to_download: List of URLs to download + temp_dir_root: Temporary directory for downloads + media_per_archive: Number of items to extract per archive + downloaded_archives: Set of already-processed parquet basenames + source: Source platform (huggingface, modelscope, s3) + hf_token: HuggingFace API token + seed: Random seed for reproducibility + + Yields: + Sample dictionaries with media and metadata + """ + parquet_urls = [url for url in to_download if ".parquet" in url] + + if not parquet_urls: + logger.warning(f"No parquet files found in download list for {dataset.name}") + return + + # Track parquet files we've already processed + if downloaded_archives is not None: + parquets_to_download = [ + url + for url in parquet_urls + if os.path.basename(url) not in downloaded_archives + ] + + skipped = len(parquet_urls) - len(parquets_to_download) + if skipped > 0: + logger.info( + f"Skipping {skipped} already-processed parquet files, " + f"downloading {len(parquets_to_download)} new parquet files" + ) + parquet_urls = parquets_to_download + + if not parquet_urls: + logger.info("No new parquet files to download") + return + + logger.info(f"Downloading {len(parquet_urls)} parquet metadata files") + + max_download_workers = min(10, len(parquet_urls)) + downloaded_parquets = download_files( + parquet_urls, + temp_dir_root, + chunk_size=8192, + max_workers=max_download_workers, + hf_token=hf_token, + ) + + if not downloaded_parquets: + logger.warning(f"Failed to download parquet files for {dataset.name}") + return + + logger.info( + f"Processing {len(downloaded_parquets)} parquet files (per-parquet download & extract)" + ) + + # Build reverse mapping: downloaded filename -> original URL. + # download_single_file() appends a URL hash to filenames (e.g., "file_abc12345.parquet"), + # but downloaded_archives.json must store original basenames (e.g., "file.parquet") + # to match HuggingFace file listings for cache completeness checks. + _downloaded_name_to_url = {} + for url in parquet_urls: + expected_name = _get_expected_download_filename(url) + _downloaded_name_to_url[expected_name] = url + + total_samples = 0 + total_archives = 0 + + # Process each parquet independently: extract metadata ā†’ download its tars ā†’ + # extract samples ā†’ mark parquet done. This ensures a parquet is only marked + # as processed after ALL of its tar archives have been successfully downloaded + # and extracted, and preserves progress incrementally across timeouts. + for parquet_path in downloaded_parquets: + if not parquet_path or not parquet_path.exists(): + continue + + try: + # Resolve original URL basename (not hash-appended) for archive tracking + original_url = _downloaded_name_to_url.get(parquet_path.name) + original_basename = os.path.basename(original_url) if original_url else parquet_path.name + + iso_week = extract_iso_week_from_path(original_url or str(parquet_path)) + + # Step 1: Extract archive filenames referenced by this parquet + archive_filenames = _extract_unique_archive_filenames(parquet_path) + if not archive_filenames: + logger.warning(f"No archive filenames found in {parquet_path}") + # Still mark as done - the parquet was processed, it just had no archives + if downloaded_archives is not None: + downloaded_archives.add(original_basename) + continue + + archive_to_week = {} + for archive_filename in archive_filenames: + if iso_week: + archive_to_week[archive_filename] = iso_week + + # Step 2: Build metadata map from this parquet + metadata_map = _build_parquet_metadata_map(parquet_path) + if iso_week: + for key in metadata_map.keys(): + if key in metadata_map: + metadata_map[key]["iso_week"] = iso_week + + # Step 3: Download this parquet's tar archives + archive_paths = [ + f"archives/{week}/{basename}" + for basename, week in archive_to_week.items() + ] + archive_urls = _get_download_urls(dataset.path, archive_paths, source) + + if not archive_urls: + logger.warning(f"No archive URLs resolved for {original_basename}") + continue + + max_workers = min(10, len(archive_urls)) + downloaded_archive_files = download_files( + archive_urls, + temp_dir_root, + chunk_size=8192, + max_workers=max_workers, + hf_token=hf_token, + ) + + if not downloaded_archive_files: + logger.warning( + f"Failed to download tar archives for parquet {original_basename}" + ) + continue + + # Step 4: Extract samples from this parquet's tar archives + parquet_samples = 0 + for archive_path in downloaded_archive_files: + if not archive_path or not archive_path.exists(): + continue + + try: + tar_iso_week = extract_iso_week_from_path(str(archive_path)) + + for sample in _process_tar_with_metadata( + archive_path, + dataset, + metadata_map, + media_per_archive, + tar_iso_week, + seed, + ): + parquet_samples += 1 + total_samples += 1 + yield sample + except (OSError, pyarrow.ArrowInvalid) as e: + logger.warning(f"Failed to process archive {archive_path}: {e}") + continue + finally: + try: + if archive_path.exists(): + archive_path.unlink() + except OSError: + pass + + total_archives += len(downloaded_archive_files) + + # Step 5: Mark parquet as done ONLY after all its tars are processed. + # This is the last step so if anything above fails or times out, + # the parquet will be retried on the next run. + if downloaded_archives is not None: + downloaded_archives.add(original_basename) + logger.info( + f"āœ… Parquet {original_basename}: {parquet_samples} samples " + f"from {len(downloaded_archive_files)} archives" + ) + except (OSError, pyarrow.ArrowInvalid) as e: + logger.warning(f"Failed to process parquet {parquet_path}: {e}") + continue + finally: + # Clean up this parquet file immediately + try: + if parquet_path and parquet_path.exists(): + parquet_path.unlink() + except (OSError, pyarrow.ArrowInvalid): + pass + + logger.info( + f"Extracted {total_samples} samples from {total_archives} archives" + ) + + + +def _download_filtered_sequential( + dataset, + filenames: List[str], + target_total: int, + temp_dir_root: Path, + source: str, + hf_token: Optional[str], + seed: Optional[int], +) -> Generator[Dict[str, Any], None, None]: + """Download parquet shards one at a time, stopping once target_total samples are collected. + + Filtering is handled inside _process_parquet via dataset.filter_column/filter_value, + so all media decoding follows the exact same path as regular parquet downloads. + """ + remote_paths = _get_download_urls(dataset.path, filenames, source) + if seed is not None: + random.Random(seed).shuffle(remote_paths) + else: + random.shuffle(remote_paths) + + collected = 0 + for url in remote_paths: + if collected >= target_total: + break + + logger.info( + f"Filtered download [{dataset.name}]: {collected}/{target_total} collected, " + f"fetching next shard" + ) + downloaded = download_files([url], temp_dir_root, hf_token=hf_token) + if not downloaded: + continue + + parquet_path = downloaded[0] + try: + remaining = target_total - collected + for sample in yield_media_from_source( + parquet_path, dataset, remaining, iso_week=None, + hf_token=hf_token, seed=seed, + ): + yield sample + collected += 1 + if collected >= target_total: + break + finally: + try: + parquet_path.unlink() + except (OSError, pyarrow.ArrowInvalid): + pass + + logger.info( + f"Filtered download complete: {collected}/{target_total} samples for {dataset.name}" + ) + + + diff --git a/src/gasbench/dataset/download/extract.py b/src/gasbench/dataset/download/extract.py new file mode 100644 index 0000000..8fbeee0 --- /dev/null +++ b/src/gasbench/dataset/download/extract.py @@ -0,0 +1,667 @@ +"""Extract and decode media samples from parquet, tar, zip, and raw sources.""" + +import base64 +import io +import json +import random +import re +import tarfile +from contextlib import closing +from io import BytesIO +from pathlib import Path +from typing import Any, Dict, Generator, Optional, Tuple +from zipfile import ZipFile +from datetime import datetime + +import soundfile as sf +import pandas as pd +import pyarrow +import pyarrow.parquet as pq +import numpy as np +from PIL import Image + +from ...logger import get_logger +from ..utils.metadata_utils import create_sample, extract_row_metadata + +from .constants import ( + IMAGE_FILE_EXTENSIONS, + VIDEO_FILE_EXTENSIONS, + AUDIO_FILE_EXTENSIONS, + _is_parquet_file, + _is_zip_file, + _is_tar_file, +) + +logger = get_logger(__name__) + + +def yield_media_from_source( + source_path: Path, + dataset, # BenchmarkDatasetConfig + num_items: int, + iso_week: Optional[str] = None, + hf_token: Optional[str] = None, + seed: Optional[int] = None, +) -> Generator[Dict[str, Any], None, None]: + """ + Unified media extractor for parquet, zip, tar sources, frame directories, and raw media files. + + Returns complete sample dictionaries ready for processing functions. + Samples include: image/video_bytes, media_type, dataset_name, dataset_path, etc. + """ + try: + if source_path.is_dir(): + yield from _process_frame_directory(source_path, dataset, iso_week) + return + + filename = str(source_path.name).lower() + + if _is_parquet_file(filename): + yield from _process_parquet(source_path, dataset, num_items, iso_week, seed) + return + + if _is_zip_file(filename) or _is_tar_file(filename): + yield from _process_zip_or_tar( + source_path, + dataset, + num_items, + iso_week, + seed, + ) + return + + if any( + filename.endswith(ext) + for ext in (IMAGE_FILE_EXTENSIONS | VIDEO_FILE_EXTENSIONS | AUDIO_FILE_EXTENSIONS) + ): + yield from _process_raw(source_path, dataset, iso_week) + return + + logger.warning(f"Unsupported source format for {source_path}") + return + except Exception as e: + logger.warning(f"Error in yield_media_from_source for {source_path}: {e}") + return + + + +def _extract_unique_archive_filenames(parquet_path: Path) -> set: + """Extract unique archive filenames from gasstation parquet metadata file. + + Args: + parquet_path: Path to the parquet metadata file + + Returns: + Set of unique archive filenames referenced in the parquet + """ + try: + table = pq.read_table(parquet_path) + df = table.to_pandas() + + if "archive_filename" not in df.columns: + logger.warning(f"No archive_filename column in {parquet_path}") + return set() + + archive_filenames = set(df["archive_filename"].dropna().unique()) + return archive_filenames + except (OSError, pyarrow.ArrowInvalid) as e: + logger.warning(f"Failed to extract archive filenames from {parquet_path}: {e}") + return set() + + + +def _build_parquet_metadata_map( + parquet_path: Path, +) -> Dict[Tuple[str, str], Dict[str, Any]]: + """Build a mapping from (archive_filename, file_path_in_archive) to metadata. + + Args: + parquet_path: Path to the parquet metadata file + + Returns: + Dictionary mapping (archive_filename, file_path_in_archive) tuples to metadata dicts + """ + try: + table = pq.read_table(parquet_path) + df = table.to_pandas() + + if ( + "archive_filename" not in df.columns + or "file_path_in_archive" not in df.columns + ): + logger.warning(f"Missing required columns in {parquet_path}") + return {} + + metadata_map = {} + for _, row in df.iterrows(): + archive_filename = row.get("archive_filename") + file_path = row.get("file_path_in_archive") + + if pd.isna(archive_filename) or pd.isna(file_path): + continue + + key = (str(archive_filename), str(file_path)) + metadata_map[key] = extract_row_metadata(row, "") + + return metadata_map + except (OSError, json.JSONDecodeError, FileNotFoundError) as e: + logger.warning(f"Failed to build metadata map from {parquet_path}: {e}") + return {} + + + +def _process_tar_with_metadata( + archive_path: Path, + dataset, + metadata_map: Dict[Tuple[str, str], Dict[str, Any]], + num_items: int, + iso_week: Optional[str], + seed: Optional[int], +) -> Generator[Dict[str, Any], None, None]: + """Extract media from tar archive using parquet metadata mapping. + + Args: + archive_path: Path to tar archive + dataset: BenchmarkDatasetConfig + metadata_map: Mapping from (archive_filename, file_path_in_archive) to metadata + num_items: Number of items to extract (-1 for all) + iso_week: ISO week string + seed: Random seed + + Yields: + Sample dictionaries with media and metadata + """ + archive_basename = archive_path.name + + archive_basename_normalized = re.sub(r"\.tar_[a-f0-9]{8}\.gz$", ".tar.gz", archive_basename) + + try: + with tarfile.open(archive_path, mode="r:*") as archive: + valid_exts = ( + IMAGE_FILE_EXTENSIONS + if dataset.modality == "image" + else (AUDIO_FILE_EXTENSIONS if dataset.modality == "audio" else VIDEO_FILE_EXTENSIONS) + ) + + all_members = [m for m in archive.getmembers() if m.isreg()] + candidates = [ + m + for m in all_members + if any(m.name.lower().endswith(ext) for ext in valid_exts) + and "MACOSX" not in m.name + ] + + if not candidates: + logger.warning(f"No matching media files found in {archive_path}") + return + + if num_items == -1: + selected = candidates + else: + if seed is not None: + rng = random.Random(seed) + selected = rng.sample(candidates, min(num_items, len(candidates))) + else: + selected = random.sample( + candidates, min(num_items, len(candidates)) + ) + + for member in selected: + try: + src = archive.extractfile(member) + if src is None: + continue + + with closing(src): + data_bytes = src.read() + + if dataset.modality == "image": + try: + media_obj = Image.open(BytesIO(data_bytes)) + except (OSError, FileNotFoundError) as e: + logger.warning(f"Failed to open image {member.name}: {e}") + continue + else: + media_obj = data_bytes + + sample = create_sample(dataset, media_obj, archive_path, iso_week) + sample["archive_filename"] = archive_basename_normalized + sample["member_path"] = member.name + + metadata_key = (archive_basename_normalized, member.name) + if metadata_key in metadata_map: + meta = metadata_map[metadata_key] + for k, v in meta.items(): + if k not in sample: + sample[k] = v + + yield sample + + except Exception as e: + logger.warning( + f"Error extracting {member.name} from {archive_path}: {e}" + ) + continue + + except Exception as e: + logger.warning(f"Error opening tar archive {archive_path}: {e}") + return + + + +def _process_zip_or_tar( + source_path: Path, + dataset, + num_items: int, + iso_week: Optional[str] = None, + seed: Optional[int] = None, +): + """Extract media from zip/tar archives (non-gasstation datasets).""" + filename = str(source_path.name).lower() + is_zip = _is_zip_file(filename) + try: + cm = ZipFile(source_path) if is_zip else tarfile.open(source_path, mode="r:*") + with cm as archive: + if dataset.modality == "image": + valid_exts = IMAGE_FILE_EXTENSIONS + elif dataset.modality == "audio": + valid_exts = AUDIO_FILE_EXTENSIONS + else: + valid_exts = VIDEO_FILE_EXTENSIONS + + if is_zip: + list_entries = archive.namelist() + + def get_name(e): + return e + + def open_entry(e): + return archive.open(e) + + else: + list_entries = [m for m in archive.getmembers() if m.isreg()] + + def get_name(m): + return m.name + + def open_entry(m): + return archive.extractfile(m) + + candidates = [ + e + for e in list_entries + if any(get_name(e).lower().endswith(ext) for ext in valid_exts) + and "MACOSX" not in get_name(e) + ] + if not candidates: + logger.warning(f"No matching files found in {source_path}") + return + + if num_items == -1: + selected = candidates + else: + if seed is not None: + rng = random.Random(seed) + selected = rng.sample(candidates, min(num_items, len(candidates))) + else: + selected = random.sample( + candidates, min(num_items, len(candidates)) + ) + + for entry in selected: + try: + src = open_entry(entry) + if src is None: + continue + with closing(src): + data_bytes = src.read() + + if dataset.modality == "image": + try: + media_obj = Image.open(BytesIO(data_bytes)) + except (OSError, FileNotFoundError): + logger.warning( + f"Failed to open image {get_name(entry)} from {source_path}" + ) + continue + else: + media_obj = data_bytes + + sample = create_sample(dataset, media_obj, source_path, iso_week) + # Normalize hash-appended filename back to original basename + # (download_single_file appends _ before extension) + _normalized = re.sub(r"\.tar_[a-f0-9]{8}\.gz$", ".tar.gz", source_path.name) + if _normalized == source_path.name: + _normalized = re.sub(r"_[a-f0-9]{8}(\.\w+)$", r"\1", source_path.name) + sample["archive_filename"] = _normalized + sample["member_path"] = get_name(entry) + + yield sample + except Exception as e: + logger.warning( + f"Error extracting {get_name(entry)} from {source_path}: {e}" + ) + continue + except Exception as e: + logger.warning(f"Error processing archive file {source_path}: {e}") + return + + + +def _process_raw(source_path: Path, dataset, iso_week: Optional[str] = None): + filename = str(source_path.name).lower() + try: + data_bytes = source_path.read_bytes() + if dataset.modality == "image" and any( + filename.endswith(ext) for ext in IMAGE_FILE_EXTENSIONS + ): + media_obj = Image.open(BytesIO(data_bytes)) + elif dataset.modality == "audio" and any( + filename.endswith(ext) for ext in AUDIO_FILE_EXTENSIONS + ): + media_obj = data_bytes + elif dataset.modality == "video" and any( + filename.endswith(ext) for ext in VIDEO_FILE_EXTENSIONS + ): + media_obj = data_bytes + else: + logger.warning( + f"Direct file {source_path} does not match modality {dataset.modality}" + ) + return + + yield create_sample(dataset, media_obj, source_path, iso_week) + except Exception as e: + logger.warning(f"Error reading direct file {source_path}: {e}") + return + + + +def _clean_to_json_serializable(value: Any) -> Any: + """Convert arbitrary values to JSON-serializable equivalents.""" + try: + if value is None or isinstance(value, (str, int, bool)): + return value + + if isinstance(value, float): + if np.isnan(value) or np.isinf(value): + return None + return value + + if isinstance(value, (np.integer,)): + return int(value) + if isinstance(value, (np.floating,)): + f = float(value) + if np.isnan(f) or np.isinf(f): + return None + return f + if isinstance(value, (np.bool_,)): + return bool(value) + + if isinstance(value, (bytes, bytearray)): + try: + return base64.b64encode(bytes(value)).decode("ascii") + except Exception: + return None + + if isinstance(value, datetime): + return value.isoformat() + + if isinstance(value, (np.ndarray, list, tuple, set)): + return [ + _clean_to_json_serializable(v) for v in (value.tolist() if isinstance(value, np.ndarray) else list(value)) + ] + + if isinstance(value, dict): + return {str(k): _clean_to_json_serializable(v) for k, v in value.items()} + + try: + json.dumps(value) + return value + except (TypeError, ValueError): + return str(value) + except (TypeError, ValueError): + return None + + + +def _extract_parquet_row_metadata(row: Any, media_col: str) -> Dict[str, Any]: + """Extract non-media columns from a pandas Series row and clean to JSON-serializable dict.""" + metadata: Dict[str, Any] = {} + try: + for col, val in row.items(): + if str(col) == str(media_col): + continue + cleaned = _clean_to_json_serializable(val) + col_str = str(col) + metadata[col_str] = cleaned + + col_lower = col_str.lower() + if "hotkey" in col_lower and "generator_hotkey" not in metadata: + metadata["generator_hotkey"] = cleaned + if "uid" in col_lower and "generator_uid" not in metadata: + metadata["generator_uid"] = cleaned + except (OSError, json.JSONDecodeError, FileNotFoundError): + pass + return metadata + + + +def _process_parquet( + source_path: Path, + dataset, + num_items: int, + iso_week: Optional[str] = None, + seed: Optional[int] = None +): + """Process parquet files with embedded media bytes (e.g., PICA-100K). + + Supports datasets with single or multiple image columns. + For datasets with multiple columns (e.g., src_img and tgt_img), yields one sample per column. + """ + try: + table = pq.read_table(source_path) + df = table.to_pandas() + + filter_column = getattr(dataset, "filter_column", None) + filter_value = getattr(dataset, "filter_value", None) + if filter_column and filter_value: + if filter_column not in df.columns: + logger.warning(f"filter_column '{filter_column}' not in {source_path.name}; skipping") + return + df = df[df[filter_column] == filter_value] + if df.empty: + logger.info(f"No rows matching {filter_column}=={filter_value} in {source_path.name}") + return + logger.info(f" {source_path.name}: {len(df)} rows after filter ({filter_column}=={filter_value})") + + if num_items == -1: + sample_df = df + else: + sample_df = df.sample(n=min(num_items, len(df)), random_state=seed) + + data_columns = getattr(dataset, 'data_columns', None) + if data_columns: + media_cols = [c for c in data_columns if c in sample_df.columns] + if not media_cols: + logger.warning( + f"Specified data_columns {data_columns} not found in {source_path}" + ) + return + elif dataset.modality == "image": + media_col = ( + next((c for c in sample_df.columns if c.lower() == "image"), None) + or next((c for c in sample_df.columns if "image" in c.lower() and "_id" not in c.lower()), None) + or next((c for c in sample_df.columns if "image" in c.lower()), None) + ) + media_cols = [media_col] if media_col else [] + elif dataset.modality == "audio": + candidates = ["audio", "bytes", "content", "data", "wav", "mp3"] + exact = [c for c in sample_df.columns if c.lower() == "audio"] + if exact: + media_cols = exact + else: + media_cols = [ + c + for c in sample_df.columns + if "_id" not in c.lower() + and any(k in c.lower() for k in candidates) + ] + else: + candidates = ["video", "bytes", "content", "data"] + media_col = ( + next((c for c in sample_df.columns if c.lower() in candidates), None) + or next((c for c in sample_df.columns if any(k in c.lower() for k in candidates) and "_id" not in c.lower()), None) + or next((c for c in sample_df.columns if any(k in c.lower() for k in candidates)), None) + ) + media_cols = [media_col] if media_col else [] + + if not media_cols: + logger.warning( + f"No media column found in {source_path} for modality {dataset.modality}" + ) + return + + for _, row in sample_df.iterrows(): + for col in media_cols: + try: + media_data = row[col] + audio_sampling_rate = None # Track sampling rate for audio dicts + + if isinstance(media_data, dict): + # For audio, try to extract sampling_rate before getting the array + if dataset.modality == "audio": + audio_sampling_rate = media_data.get("sampling_rate") or media_data.get("sample_rate") + + key = next( + ( + k + for k in media_data + if any( + s in k.lower() + for s in ["bytes", "image", "video", "audio", "array", "data", "content"] + ) + ), + None, + ) + if key: + media_data = media_data[key] + else: + logger.warning(f"No valid key found in dict media_data for {source_path}: {list(media_data.keys())}") + continue + + if dataset.modality == "image": + if media_data is None or isinstance(media_data, (int, float)): + continue + + try: + img = Image.open(BytesIO(media_data)) + except (OSError, FileNotFoundError): + if isinstance(media_data, str): + media_data = base64.b64decode(media_data) + img = Image.open(BytesIO(media_data)) + sample = create_sample(dataset, img, source_path, iso_week) + else: + if media_data is None or isinstance(media_data, (int, float)): + continue + + # Handle audio array format (common in HF audio datasets) + if dataset.modality == "audio" and isinstance(media_data, (list, np.ndarray)): + # Audio data is a numpy array - convert to WAV bytes + try: + audio_array = np.array(media_data, dtype=np.float32) + if audio_array.ndim == 1: + audio_array = audio_array.reshape(-1, 1) + sr = audio_sampling_rate if audio_sampling_rate else 16000 + buffer = io.BytesIO() + sf.write(buffer, audio_array, sr, format='WAV') + media_data = buffer.getvalue() + except (ValueError, TypeError, OSError) as e: + logger.warning(f"Failed to convert audio array to bytes: {e}") + continue + elif not isinstance(media_data, (bytes, bytearray)): + if isinstance(media_data, str): + media_data = base64.b64decode(media_data) + else: + continue + sample = create_sample(dataset, bytes(media_data), source_path, iso_week) + + row_metadata = _extract_parquet_row_metadata(row, col) + for k, v in row_metadata.items(): + if k not in sample: + sample[k] = v + + if len(media_cols) > 1: + sample["source_column"] = col + + yield sample + except Exception as e: + logger.warning(f"Failed to extract row from {source_path} (column {col}): {e}") + continue + except Exception as e: + logger.warning(f"Error processing parquet file {source_path}: {e}") + return + + + +def _process_frame_directory( + frame_dir: Path, dataset, iso_week: Optional[str] = None +) -> Generator[Dict[str, Any], None, None]: + """Process a directory containing pre-extracted video frames. + + Args: + frame_dir: Directory containing frame files (e.g., bm-videos/dfb/DFDC/test/frames/amwhgrjvkw/) + dataset: BenchmarkDatasetConfig + iso_week: Optional ISO week string + + Yields: + Sample dict with 'video_frames' key containing list of frame paths + """ + try: + if not frame_dir.is_dir(): + logger.warning(f"{frame_dir} is not a directory") + return + + # Get all image files in the directory + frame_files = [] + for ext in IMAGE_FILE_EXTENSIONS: + frame_files.extend(sorted(frame_dir.glob(f"*{ext}"))) + frame_files.extend(sorted(frame_dir.glob(f"*{ext.upper()}"))) + + # Sort frames numerically if they have numeric names + def extract_number(filepath): + """Extract numeric part from filename for sorting (e.g., '000.png' -> 0)""" + name = filepath.stem + match = re.match(r"(\d+)", name) + if match: + return int(match.group(1)) + return name + + frame_files = sorted(set(frame_files), key=extract_number) + + if not frame_files: + logger.warning(f"No frame files found in {frame_dir}") + return + + logger.debug(f"Found {len(frame_files)} frames in {frame_dir.name}") + + # Create sample with frame paths + sample = { + "video_frames": frame_files, + "media_type": dataset.media_type, + "dataset_name": dataset.name, + "dataset_path": dataset.path, + "source_file": frame_dir.name, + } + + if iso_week: + sample["iso_week"] = iso_week + + yield sample + + except Exception as e: + logger.warning(f"Error processing frame directory {frame_dir}: {e}") + return + + + diff --git a/src/gasbench/dataset/download/fetch.py b/src/gasbench/dataset/download/fetch.py new file mode 100644 index 0000000..ad5a1a5 --- /dev/null +++ b/src/gasbench/dataset/download/fetch.py @@ -0,0 +1,247 @@ +"""Low-level file downloading (HTTP/HF streaming) and output-name helpers.""" + +import hashlib +import os +import time +import traceback +from pathlib import Path +from typing import Generator, List, Optional + +import requests + +from ...logger import get_logger +from ..utils.s3_utils import download_s3_file + +logger = get_logger(__name__) + +def _stream_downloads( + urls: List[str], + output_dir: Path, + chunk_size: int = 8192, + max_workers: int = 10, + hf_token: Optional[str] = None, +) -> Generator[Optional[Path], None, None]: + """Like download_files but yields each path as soon as its download completes.""" + from concurrent.futures import ThreadPoolExecutor, as_completed + + output_dir.mkdir(parents=True, exist_ok=True) + if not urls: + return + + def download_url(url): + try: + return download_single_file(url, output_dir, chunk_size, hf_token) + except (requests.RequestException, OSError) as e: + logger.error(f"Error downloading {url}: {e}") + return None + + workers = min(max_workers, len(urls)) + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = {executor.submit(download_url, url): url for url in urls} + for future in as_completed(futures): + yield future.result() + + + +def download_files( + urls: List[str], + output_dir: Path, + chunk_size: int = 8192, + max_workers: int = 10, + hf_token: Optional[str] = None, +) -> List[Path]: + """Download multiple files in parallel. + + Args: + urls: List of URLs to download + output_dir: Directory to save the files + chunk_size: Size of chunks to download at a time + max_workers: Maximum number of parallel downloads (default: 10) + hf_token: Hugging Face API token for private datasets + + Returns: + List of successfully downloaded file paths + """ + from concurrent.futures import ThreadPoolExecutor, as_completed + + output_dir.mkdir(parents=True, exist_ok=True) + downloaded_files = [] + + def download_url(url): + try: + return download_single_file(url, output_dir, chunk_size, hf_token) + except (requests.RequestException, OSError) as e: + logger.error(f"Error downloading {url}: {e}") + return None + + max_workers = min(max_workers, len(urls)) if urls else 1 + + logger.debug(f"Downloading {len(urls)} files with {max_workers} parallel workers") + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(download_url, url): url for url in urls} + + for future in as_completed(futures): + result = future.result() + if result: + downloaded_files.append(result) + + return downloaded_files + + + +def download_single_file( + url: str, output_dir: Path, chunk_size: int, hf_token: Optional[str] = None +) -> Optional[Path]: + """Download a single file synchronously with resume and retry support + + Args: + url: URL to download (supports http/https and s3: schemes) + output_dir: Directory to save the file + chunk_size: Size of chunks to download at a time + hf_token: Hugging Face API token for private datasets + + Returns: + Path to the downloaded file, or None if failed + """ + try: + # Handle S3 URLs (format: s3:bucket-name/key/path or s3:bucket-name/frame/directory/) + if url.startswith("s3:"): + from .utils.s3_utils import download_s3_frame_directory + + s3_path = url[3:] # Remove "s3:" prefix + parts = s3_path.split("/", 1) + if len(parts) != 2: + logger.error( + f"Invalid S3 URL format: {url}. Expected: s3:bucket/key/path" + ) + return None + + bucket = parts[0] + key = parts[1] + + # Check if this is a frame directory (no file extension or ends with /) + if key.endswith("/") or ("." not in os.path.basename(key)): + # This is a frame directory - download all frames + return download_s3_frame_directory(bucket, key, output_dir) + else: + # This is a single file - use unique filename to avoid collisions + url_hash = hashlib.md5(key.encode()).hexdigest()[:8] + base_filename = os.path.basename(key) + name, ext = os.path.splitext(base_filename) + filename = f"{name}_{url_hash}{ext}" + filepath = output_dir / filename + return download_s3_file(bucket, key, filepath) + + # Handle regular HTTP/HTTPS URLs + # Use hash of full URL path to avoid filename collisions when multiple + # files from different directories have the same basename + url_hash = hashlib.md5(url.encode()).hexdigest()[:8] + base_filename = os.path.basename(url) + name, ext = os.path.splitext(base_filename) + filename = f"{name}_{url_hash}{ext}" + filepath = output_dir / filename + partial_filepath = Path(str(filepath) + ".partial") + + logger.debug(f"Downloading {url}") + + max_retries = 5 + effective_chunk_size = max(chunk_size, 1024 * 1024) + + for attempt in range(max_retries): + try: + headers = {} + if hf_token and "huggingface.co" in url: + headers["Authorization"] = f"Bearer {hf_token}" + + downloaded_size = 0 + if partial_filepath.exists(): + downloaded_size = partial_filepath.stat().st_size + headers["Range"] = f"bytes={downloaded_size}-" + logger.info(f"Resuming download from {downloaded_size} bytes") + + response = requests.get( + url, stream=True, timeout=(30, 300), headers=headers + ) + + if response.status_code == 416: + logger.info("Server returned 416 (range not satisfiable), restarting download") + partial_filepath.unlink(missing_ok=True) + downloaded_size = 0 + del headers["Range"] + response = requests.get( + url, stream=True, timeout=(30, 300), headers=headers + ) + + if response.status_code not in (200, 206): + logger.error(f"Failed to download {url}: Status {response.status_code}") + return None + + if response.status_code == 200 and downloaded_size > 0: + logger.info("Server doesn't support resume, restarting download") + partial_filepath.unlink(missing_ok=True) + downloaded_size = 0 + + total_size = int(response.headers.get("content-length", 0)) + if response.status_code == 200: + expected_total = total_size + else: + expected_total = downloaded_size + total_size + + mode = "ab" if downloaded_size > 0 else "wb" + with open(partial_filepath, mode) as f: + for chunk in response.iter_content(chunk_size=effective_chunk_size): + if chunk: + f.write(chunk) + + actual_size = partial_filepath.stat().st_size + if expected_total > 0 and actual_size != expected_total: + raise IOError( + f"Download incomplete: expected {expected_total} bytes, got {actual_size} bytes" + ) + + partial_filepath.rename(filepath) + return filepath + + except (requests.exceptions.RequestException, IOError) as e: + wait_time = min(30 * (2 ** attempt), 300) + logger.warning( + f"Download attempt {attempt + 1}/{max_retries} failed: {e}. " + f"Retrying in {wait_time}s..." + ) + if attempt < max_retries - 1: + time.sleep(wait_time) + else: + raise + + return None + + except Exception as e: + logger.error(f"Error downloading {url}: {str(e)}") + logger.error(traceback.format_exc()) + return None + + + +def _get_expected_download_filename(url: str) -> str: + """Compute the filename that download_single_file() would produce for a URL. + + download_single_file() appends a URL hash to filenames to avoid collisions + (e.g., "file.parquet" becomes "file_abc12345.parquet"). This function mirrors + that naming convention to enable reverse-mapping from downloaded filenames + back to original URLs. + """ + if url.startswith("s3:"): + s3_path = url[3:] + parts = s3_path.split("/", 1) + key = parts[1] if len(parts) == 2 else s3_path + url_hash = hashlib.md5(key.encode()).hexdigest()[:8] + base_filename = os.path.basename(key) + else: + url_hash = hashlib.md5(url.encode()).hexdigest()[:8] + base_filename = os.path.basename(url) + name, ext = os.path.splitext(base_filename) + return f"{name}_{url_hash}{ext}" + + + diff --git a/src/gasbench/dataset/download/listing.py b/src/gasbench/dataset/download/listing.py new file mode 100644 index 0000000..cd27c55 --- /dev/null +++ b/src/gasbench/dataset/download/listing.py @@ -0,0 +1,296 @@ +"""Remote dataset file listing and download-URL resolution (HF/ModelScope/S3).""" + +import random +from urllib.parse import quote +from typing import List, Optional + +import huggingface_hub as hf_hub +from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError +from modelscope.hub.api import HubApi as MSHubApi + +from ...logger import get_logger +from ..utils.s3_utils import list_s3_files, _get_s3_urls +from ..utils.gasstation_utils import ( + filter_files_by_current_week, + filter_files_by_recent_weeks, +) + +from .constants import DatasetAccessError, MAX_FILES_DEFAULT + +logger = get_logger(__name__) + + +def _select_files_to_download( + urls: List[str], + count: int, + prioritize_recent: bool = False, + seed: Optional[int] = None, +) -> List[str]: + """Select files to download. + + Args: + urls: List of file URLs to select from + count: Number of files to select (-1 = all files) + prioritize_recent: If True, sort by filename (descending) to get newest first. + Useful for gasstation datasets to get freshest data. + seed: Random seed for reproducible sampling (non-gasstation only) + + Returns: + List of selected URLs + """ + if count == -1: + return urls + if count <= 0: + return [] + + if prioritize_recent: + sorted_urls = sorted(urls, reverse=True) + return sorted_urls[: min(count, len(urls))] + + if seed is not None: + rng = random.Random(seed) + return rng.sample(urls, min(count, len(urls))) + return random.sample(urls, min(count, len(urls))) + + + +def _list_remote_dataset_files( + dataset_path: str, + source_format: str = ".parquet", + current_week_only: bool = False, + num_weeks: int = None, + target_week: str = None, + include_paths: Optional[List[str]] = None, + exclude_paths: Optional[List[str]] = None, + source: str = "huggingface", + hf_token: Optional[str] = None, + max_files: Optional[int] = None, +) -> List[str]: + """List available files in a dataset, filtered by source_format and path patterns. + + Supports single extensions (e.g., .parquet, .zip) and tar variants (.tar, .tar.gz, .tgz). + For gasstation datasets with current_week_only=True, filters to only current ISO week's data. + For gasstation datasets with num_weeks set, filters to only the N most recent weeks. + + Args: + dataset_path: Dataset repository path (org/dataset-name) + source_format: File extension to filter by + current_week_only: For gasstation datasets, only get current week + num_weeks: For gasstation datasets, get N most recent weeks + target_week: For gasstation datasets, get specific week + include_paths: Only include files containing one of these path segments + exclude_paths: Exclude files containing any of these path segments + source: Source platform ("huggingface" or "modelscope") + max_files: Stop after collecting this many files (None uses MAX_FILES_DEFAULT) + """ + if max_files is None: + max_files = MAX_FILES_DEFAULT + # Don't add "." prefix for special formats like "frames" + if source_format != "frames" and not source_format.startswith("."): + source_format = "." + source_format + + if source_format in [".tar", ".tar.gz", ".tgz"]: + source_format = [".tar", ".tar.gz", ".tgz"] + + # For gasstation datasets with week filtering, we must fetch ALL files before + # applying max_files. Otherwise early termination picks up files from older + # weeks (alphabetically first) and the subsequent week filter discards them all, + # resulting in zero files. + is_gasstation = "gasstation" in dataset_path.lower() + needs_week_filter = is_gasstation and (target_week or num_weeks or current_week_only) + listing_max = None if needs_week_filter else max_files + + if source == "modelscope": + files = list_modelscope_files(repo_id=dataset_path, extension=source_format) + if include_paths: + files = [f for f in files if any(path_seg in f for path_seg in include_paths)] + if exclude_paths: + files = [f for f in files if not any(path_seg in f for path_seg in exclude_paths)] + if not needs_week_filter: + files = files[:max_files] + elif source == "s3": + files = list_s3_files(path=dataset_path, extension=source_format) + if include_paths: + files = [f for f in files if any(path_seg in f for path_seg in include_paths)] + if exclude_paths: + files = [f for f in files if not any(path_seg in f for path_seg in exclude_paths)] + if not needs_week_filter: + files = files[:max_files] + else: # hf - supports early termination natively + files = list_hf_files( + repo_id=dataset_path, + extension=source_format, + token=hf_token, + max_files=listing_max, + include_paths=include_paths, + exclude_paths=exclude_paths, + ) + + if is_gasstation: + if target_week: + files = [f for f in files if target_week in f] + logger.info( + f"Filtered to week {target_week} for {dataset_path}: {len(files)} files" + ) + elif num_weeks: + files = filter_files_by_recent_weeks(files, num_weeks) + logger.info( + f"Filtered to last {num_weeks} weeks for {dataset_path}: {len(files)} files" + ) + elif current_week_only: + files = filter_files_by_current_week(files) + logger.info( + f"Filtered to current week files for {dataset_path}: {len(files)} files" + ) + + # Apply max_files limit AFTER week filtering + if needs_week_filter and max_files: + files = files[:max_files] + + return files + + + +def _get_download_urls( + dataset_path: str, filenames: List[str], source: str = "huggingface" +) -> List[str]: + """Get download URLs for data files from the specified source. + + Args: + dataset_path: Repository path (org/dataset-name) or S3 path (bucket/prefix) + filenames: List of files to download + source: Source platform ("huggingface", "modelscope", or "s3") + + Returns: + List of download URLs + """ + if source == "modelscope": + return _get_modelscope_urls(dataset_path, filenames) + elif source == "s3": + return _get_s3_urls(dataset_path, filenames) + else: + return _get_huggingface_urls(dataset_path, filenames) + + + +def _get_huggingface_urls(dataset_path: str, filenames: List[str]) -> List[str]: + return [ + f"https://huggingface.co/datasets/{dataset_path}/resolve/main/{quote(f, safe='/')}" + for f in filenames + ] + + + +def _get_modelscope_urls(dataset_path: str, filenames: List[str]) -> List[str]: + return [ + f"https://www.modelscope.cn/api/v1/datasets/{dataset_path}/repo?Revision=master&FilePath={f}" + for f in filenames + ] + + + +def list_hf_files( + repo_id, + repo_type="dataset", + extension=None, + token=None, + max_files=None, + include_paths=None, + exclude_paths=None, +): + """List files from a Hugging Face repository with early termination support. + + Args: + repo_id: Repository ID + repo_type: Type of repository ('dataset', 'model', etc.) + extension: Filter files by extension + token: Hugging Face API token for private datasets + max_files: Stop after collecting this many matching files (None = no limit) + include_paths: Only include files containing one of these path segments + exclude_paths: Exclude files containing any of these path segments + + Returns: + List of files in the repository + + Raises: + DatasetAccessError: If the repository is gated/private without access, or not found + """ + files = [] + if extension: + if isinstance(extension, (list, tuple, set)): + exts = tuple(extension) + else: + exts = (extension,) + else: + exts = None + + try: + for f in hf_hub.list_repo_files(repo_id=repo_id, repo_type=repo_type, token=token): + if exts and not f.endswith(exts): + continue + if include_paths and not any(path_seg in f for path_seg in include_paths): + continue + if exclude_paths and any(path_seg in f for path_seg in exclude_paths): + continue + files.append(f) + if max_files and len(files) >= max_files: + logger.info(f"Early termination: collected {max_files} files from {repo_id}") + break + except GatedRepoError: + raise DatasetAccessError( + f"Dataset {repo_id} is gated and requires access approval. " + "Visit the dataset page on HuggingFace to request access." + ) + except RepositoryNotFoundError: + raise DatasetAccessError( + f"Dataset {repo_id} not found. It may not exist or may be private." + ) + except Exception as e: + logger.error(f"Failed to list files of type {extension} in {repo_id}: {e}") + return files + + + +def list_modelscope_files(repo_id, extension=None): + """List files from a ModelScope dataset repository. + + Args: + repo_id: Repository ID (format: org/dataset-name) + extension: Filter files by extension(s) + + Returns: + List of files in the repository + """ + files = [] + try: + api = MSHubApi() + file_info = api.get_dataset_files(repo_id, revision="master", recursive=True) + + # Extract file paths + if isinstance(file_info, list): + files = [ + f["Path"] if isinstance(f, dict) and "Path" in f else str(f) + for f in file_info + ] + elif isinstance(file_info, dict): + files = [f for f in file_info.keys()] if file_info else [] + + if extension and files: + if isinstance(extension, (list, tuple, set)): + exts = tuple(extension) + files = [f for f in files if f.endswith(exts)] + else: + files = [f for f in files if f.endswith(extension)] + + except Exception as e: + logger.error( + f"Failed to list files of type {extension} in ModelScope repo {repo_id}: {e}" + ) + logger.error( + "Make sure 'modelscope' package is installed: pip install modelscope" + ) + + return files + + + diff --git a/src/gasbench/dataset/iterator.py b/src/gasbench/dataset/iterator.py index c69280a..bd68643 100644 --- a/src/gasbench/dataset/iterator.py +++ b/src/gasbench/dataset/iterator.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import Optional, Dict -from PIL import Image from ..logger import get_logger from .config import BenchmarkDatasetConfig diff --git a/src/gasbench/dataset/utils/gasstation_utils.py b/src/gasbench/dataset/utils/gasstation_utils.py index 6d89ce2..50a3a9b 100644 --- a/src/gasbench/dataset/utils/gasstation_utils.py +++ b/src/gasbench/dataset/utils/gasstation_utils.py @@ -9,7 +9,7 @@ import os import json -from typing import List, Set, Optional, Dict, Tuple +from typing import List, Set, Optional, Tuple from pathlib import Path from datetime import datetime, timedelta diff --git a/src/gasbench/download_manager.py b/src/gasbench/download_manager.py index f0ef63c..24b9891 100644 --- a/src/gasbench/download_manager.py +++ b/src/gasbench/download_manager.py @@ -149,7 +149,7 @@ async def _download_dataset(self, task: DatasetTask, progress: Progress): await asyncio.sleep(0.3) progress.remove_task(task_id) - except Exception as e: + except Exception: progress.update(task_id, description=f"[red]āŒ {task.dataset.name[:45]}") await asyncio.sleep(0.3) progress.remove_task(task_id) diff --git a/src/gasbench/processing/archive.py b/src/gasbench/processing/archive.py index fc32465..198b1ad 100644 --- a/src/gasbench/processing/archive.py +++ b/src/gasbench/processing/archive.py @@ -65,7 +65,7 @@ def get_video_path(self, sample: Dict) -> Optional[str]: contents = os.listdir(cache_dir) logger.debug(f"šŸ“ Completed archive contents: {contents}") self._logged_structure_for.add(cache_dir) - except: + except Exception: pass return None @@ -157,7 +157,7 @@ def _download_and_extract_archive( with open(completion_marker, "w") as f: f.write(f"Extracted at {time.time()}") - logger.info(f"āœ… Archive extracted successfully") + logger.info("āœ… Archive extracted successfully") except Exception as extract_error: logger.error(f"āŒ Failed to extract archive: {extract_error}") @@ -205,12 +205,12 @@ def list_directory_tree(directory, prefix=""): ) else: items.append(f"{prefix}{item}") - except: + except Exception: pass return items tree = list_directory_tree(cache_dir) - logger.debug(f"šŸ“ Archive directory structure:") + logger.debug("šŸ“ Archive directory structure:") for item in tree[:20]: # Limit to first 20 items logger.debug(f" {item}") if len(tree) > 20: diff --git a/src/gasbench/processing/media.py b/src/gasbench/processing/media.py index e0d2e89..9bb8009 100644 --- a/src/gasbench/processing/media.py +++ b/src/gasbench/processing/media.py @@ -208,11 +208,11 @@ def process_video_bytes_sample( finally: try: os.unlink(temp_video_path) - except: + except Exception: pass try: vr.close() - except: + except Exception: pass except Exception as e: @@ -262,7 +262,7 @@ def process_video_frames_sample( nparr = np.frombuffer(frame_data, np.uint8) frame_array = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if frame_array is None: - logger.warning(f"Failed to decode frame bytes") + logger.warning("Failed to decode frame bytes") continue frame_array = cv2.cvtColor(frame_array, cv2.COLOR_BGR2RGB) else: diff --git a/src/gasbench/processing/preprocess_audio.py b/src/gasbench/processing/preprocess_audio.py index 9345289..5456cba 100644 --- a/src/gasbench/processing/preprocess_audio.py +++ b/src/gasbench/processing/preprocess_audio.py @@ -24,7 +24,6 @@ import os import argparse import torch -from pathlib import Path from tqdm import tqdm from ..logger import get_logger @@ -66,7 +65,7 @@ def preprocess_dataset( samples_dir = os.path.join(dataset_dir, "samples") if not os.path.exists(samples_dir): - logger.error(f"Dataset not cached. Please run 'gasbench download --modality audio' first") + logger.error("Dataset not cached. Please run 'gasbench download --modality audio' first") return False # Check if already preprocessed @@ -145,7 +144,7 @@ def preprocess_dataset( logger.warning(f"Failed to preprocess sample {sample_idx}: {e}") error_count += 1 - logger.info(f"Preprocessing complete!") + logger.info("Preprocessing complete!") logger.info(f" Processed: {processed_count}") logger.info(f" Errors: {error_count}") logger.info(f" Location: {samples_dir}") diff --git a/src/gasbench/processing/transforms.py b/src/gasbench/processing/transforms.py index 863e72d..91573a6 100644 --- a/src/gasbench/processing/transforms.py +++ b/src/gasbench/processing/transforms.py @@ -1,13 +1,9 @@ import math import random from scipy import ndimage -import torchvision.transforms as transforms -import torchvision.transforms.functional as F import numpy as np import cv2 import os -from pathlib import Path -import time from io import BytesIO from PIL import Image import torch @@ -101,9 +97,6 @@ def apply_robustness_augmentations( in PrefetchPipeline when robustness_pass=True. Deterministic given seed so paired base/augmented samples share sample_id for degradation join. """ - if seed is not None: - np.random.seed(seed) - img = image_array.copy() if img.dtype != np.uint8: img = np.clip(img, 0, 255).astype(np.uint8) @@ -175,7 +168,6 @@ def _h264_roundtrip_ffmpeg(video_array, crf, fps): import shutil import subprocess import tempfile - import os if shutil.which("ffmpeg") is None: return None @@ -214,7 +206,6 @@ def _h264_roundtrip_cv2(video_array, crf, fps): directly, so it is approximated through VIDEOWRITER_PROP_QUALITY (a perceptual 0-100 knob that some builds ignore). Returns (T, H, W, 3) uint8 or None.""" import tempfile - import os T, H, W, C = video_array.shape # Map CRF [18, 51] ā†’ cv2 quality [100, 0] linearly (approximate only) @@ -276,9 +267,6 @@ def apply_video_robustness_augmentations( Returns the same 4-tuple as apply_random_augmentations for drop-in use in VideoPrefetchPipeline when robustness_pass=True. """ - if seed is not None: - np.random.seed(seed) - if video_array.dtype != np.uint8: video_array = np.clip(video_array, 0, 255).astype(np.uint8) @@ -367,9 +355,18 @@ def apply_random_augmentations( Raises: ValueError: If probabilities don't sum to 1.0 (within floating point precision) """ + # Use local generators seeded per-call rather than seeding the global RNG. + # The pipeline runs these in parallel worker threads; seeding the process-wide + # np.random/random state would let concurrent threads clobber each other's + # seed, breaking the determinism the seed is supposed to guarantee. + # RandomState/random.Random reproduce the legacy global-seed algorithm exactly, + # so seeded outputs (and the aug cache) are unchanged. if seed is not None: - random.seed(seed) - np.random.seed(seed) + rng = np.random.RandomState(seed) + pyrng = random.Random(seed) + else: + rng = np.random + pyrng = random if level is None: if level_probs is None: level_probs = { @@ -389,7 +386,7 @@ def apply_random_augmentations( cumsum += prob cumulative_probs[level] = cumsum - rand_val = np.random.random() + rand_val = rng.random_sample() for curr_level, cum_prob in cumulative_probs.items(): if rand_val <= cum_prob: level = curr_level @@ -397,20 +394,20 @@ def apply_random_augmentations( # determine crop scale for h and w crop_scale = (1., 1.) - if np.random.rand() < crop_prob: - crop_scale = (np.random.uniform(0.35, 0.99), np.random.uniform(0.35, 0.99)) + if rng.rand() < crop_prob: + crop_scale = (rng.uniform(0.35, 0.99), rng.uniform(0.35, 0.99)) min_scale_h = max(0.35, 224 / target_size[0]) min_scale_w = max(0.35, 224 / target_size[1]) crop_scale = (max(crop_scale[0], min_scale_h), max(crop_scale[1], min_scale_w)) if level == 0: - tforms = get_base_transforms(target_size, crop_scale) + tforms = get_base_transforms(target_size, crop_scale, rng, pyrng) elif level == 1: - tforms = get_random_augmentations(target_size, crop_scale) + tforms = get_random_augmentations(target_size, crop_scale, rng, pyrng) elif level == 2: - tforms = get_random_augmentations_medium(target_size, crop_scale) + tforms = get_random_augmentations_medium(target_size, crop_scale, rng, pyrng) else: # level == 3 - tforms = get_random_augmentations_hard(target_size, crop_scale) + tforms = get_random_augmentations_hard(target_size, crop_scale, rng, pyrng) if isinstance(inputs, tuple): transformed_A, _ = tforms(inputs[0], reuse_params=False) @@ -423,13 +420,15 @@ def apply_random_augmentations( return transformed_inputs, transformed_masks, level, tforms.params -def get_base_transforms(target_res, crop_scale): +def get_base_transforms(target_res, crop_scale, rng=None, pyrng=None): """ Get basic transforms (optional crop and resize). Args: target_res: int or tuple. Output size for resize. crop_scale: tuple. Crop scale (ignored if (1.0, 1.0)). + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). Returns: ComposeWithParams: Composed transform pipeline @@ -437,71 +436,77 @@ def get_base_transforms(target_res, crop_scale): transforms_list = [] if crop_scale != (1.0, 1.0): - transforms_list.append(RandomCropWithParams(crop_scale)) + transforms_list.append(RandomCropWithParams(crop_scale, rng=rng)) transforms_list.append(ResizeShortestEdge(target_res)) return ComposeWithParams(transforms_list) -def get_random_augmentations(target_res, crop_scale): +def get_random_augmentations(target_res, crop_scale, rng=None, pyrng=None): """ Get basic augmentations with geometric transforms. Args: target_res: int or tuple. Output size for resize. crop_scale: tuple. Crop scale. + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). Returns: ComposeWithParams: Composed transform pipeline with basic augmentations """ - base = get_base_transforms(target_res, crop_scale) + base = get_base_transforms(target_res, crop_scale, rng, pyrng) transforms_list = base.transforms + [ - RandomHorizontalFlipWithParams(), - RandomVerticalFlipWithParams(), + RandomHorizontalFlipWithParams(rng=rng), + RandomVerticalFlipWithParams(rng=rng), ] return ComposeWithParams(transforms_list) -def get_random_augmentations_medium(target_res, crop_scale): +def get_random_augmentations_medium(target_res, crop_scale, rng=None, pyrng=None): """ Get medium difficulty transforms with mild distortions. Args: target_res: int or tuple. Output size for resize. crop_scale: tuple. Crop scale. + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). Returns: ComposeWithParams: Composed transform pipeline with medium distortions """ - base = get_base_transforms(target_res, crop_scale) + base = get_base_transforms(target_res, crop_scale, rng, pyrng) transforms_list = base.transforms + [ - RandomHorizontalFlipWithParams(), - RandomVerticalFlipWithParams(), - ApplyDeeperForensicsDistortion("CS", level_min=0, level_max=1), - ApplyDeeperForensicsDistortion("CC", level_min=0, level_max=1), + RandomHorizontalFlipWithParams(rng=rng), + RandomVerticalFlipWithParams(rng=rng), + ApplyDeeperForensicsDistortion("CS", level_min=0, level_max=1, rng=rng, pyrng=pyrng), + ApplyDeeperForensicsDistortion("CC", level_min=0, level_max=1, rng=rng, pyrng=pyrng), ] return ComposeWithParams(transforms_list) -def get_random_augmentations_hard(target_res, crop_scale): +def get_random_augmentations_hard(target_res, crop_scale, rng=None, pyrng=None): """ Get hard difficulty transforms with more severe distortions. Args: target_res: int or tuple. Output size for resize. crop_scale: tuple. Crop scale. + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). Returns: ComposeWithParams: Composed transform pipeline with severe distortions """ - base = get_base_transforms(target_res, crop_scale) + base = get_base_transforms(target_res, crop_scale, rng, pyrng) transforms_list = base.transforms + [ - RandomHorizontalFlipWithParams(), - RandomVerticalFlipWithParams(), - ApplyDeeperForensicsDistortion("CS", level_min=0, level_max=2), - ApplyDeeperForensicsDistortion("CC", level_min=0, level_max=2), - ApplyDeeperForensicsDistortion("GNC", level_min=0, level_max=2), - ApplyDeeperForensicsDistortion("GB", level_min=0, level_max=2), + RandomHorizontalFlipWithParams(rng=rng), + RandomVerticalFlipWithParams(rng=rng), + ApplyDeeperForensicsDistortion("CS", level_min=0, level_max=2, rng=rng, pyrng=pyrng), + ApplyDeeperForensicsDistortion("CC", level_min=0, level_max=2, rng=rng, pyrng=pyrng), + ApplyDeeperForensicsDistortion("GNC", level_min=0, level_max=2, rng=rng, pyrng=pyrng), + ApplyDeeperForensicsDistortion("GB", level_min=0, level_max=2, rng=rng, pyrng=pyrng), ] return ComposeWithParams(transforms_list) @@ -612,7 +617,7 @@ def apply_transform(self, image, mask, transform): class ApplyDeeperForensicsDistortion: """Wrapper for applying DeeperForensics distortions.""" - def __init__(self, distortion_type, level_min=0, level_max=3): + def __init__(self, distortion_type, level_min=0, level_max=3, rng=None, pyrng=None): """ Initialize distortion transform. @@ -620,12 +625,16 @@ def __init__(self, distortion_type, level_min=0, level_max=3): distortion_type: str, type of distortion to apply level_min: int, minimum distortion level level_max: int, maximum distortion level + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). """ self.__name__ = distortion_type self.distortion_type = distortion_type self.level = None self.level_min = level_min self.level_max = level_max + self.rng = rng + self.pyrng = pyrng self.params = {} # level self.distortion_params = {} # distortion_type specific @@ -642,7 +651,8 @@ def __call__(self, img, level=None, **kwargs): np.ndarray: Distorted image """ if level is None and self.level is None: - self.level = random.randint(self.level_min, self.level_max) + pyrng = self.pyrng if self.pyrng is not None else random + self.level = pyrng.randint(self.level_min, self.level_max) self.params = {"level": self.level} elif self.level is None: self.level = level @@ -658,7 +668,11 @@ def __call__(self, img, level=None, **kwargs): else: return img - output = self.distortion_func(img, **self.distortion_params) + # Only the noise distortions consume RNG; pass the local generator to them. + if self.distortion_type in ("GNC", "BW"): + output = self.distortion_func(img, rng=self.rng, **self.distortion_params) + else: + output = self.distortion_func(img, **self.distortion_params) if isinstance(output, tuple): self.distortion_params.update(output[1]) return output[0] @@ -815,7 +829,7 @@ def color_contrast(img, param): return img.astype(np.uint8) -def block_wise(img, param): +def block_wise(img, param, rng=None): """Apply block-wise distortion by adding random gray blocks. NOTE: CURRENTLY NOT USED @@ -823,35 +837,39 @@ def block_wise(img, param): Args: img (np.ndarray): Input RGB image array of shape (H, W, 3) param (int): Number of blocks to add, scaled by image dimensions + rng: np.random.RandomState (or None for the global module). Returns: np.ndarray: Distorted RGB image array with added gray blocks """ + draw = rng if rng is not None else np.random width = 8 block = np.ones((width, width, 3)).astype(int) * 128 param = min(img.shape[0], img.shape[1]) // 256 * param for _ in range(param): - r_w = random.randint(0, img.shape[1] - 1 - width) - r_h = random.randint(0, img.shape[0] - 1 - width) + r_w = draw.randint(0, img.shape[1] - width) + r_h = draw.randint(0, img.shape[0] - width) img[r_h : r_h + width, r_w : r_w + width, :] = block return img -def gaussian_noise_color(img, param): +def gaussian_noise_color(img, param, rng=None): """Apply colored Gaussian noise in YCbCr color space. Args: img (np.ndarray): Input RGB image array of shape (H, W, 3) param (float): Variance of the Gaussian noise + rng: np.random.RandomState (or None for the global module). Returns: tuple: (distorted_image, params) distorted_image: np.ndarray, image with added color noise """ + draw = rng if rng is not None else np.random ycbcr = rgb2ycbcr(img) / 255 size_a = ycbcr.shape b = ( - ycbcr + math.sqrt(param) * np.random.randn(size_a[0], size_a[1], size_a[2]) + ycbcr + math.sqrt(param) * draw.randn(size_a[0], size_a[1], size_a[2]) ) * 255 b = ycbcr2rgb(b) return np.clip(b, 0, 255).astype(np.uint8) @@ -1077,15 +1095,17 @@ def __call__(self, img, mask=None): class RandomCropWithParams: """Randomly crop an image with optional mask-aware cropping.""" - def __init__(self, crop_scale): + def __init__(self, crop_scale, rng=None): """ Initialize random crop transform. Args: crop_scale (tuple): (height_scale, width_scale) for crop size as fraction of original + rng: np.random.RandomState (or None for the global module). """ self.params = None self.crop_scale = crop_scale + self.rng = rng def __call__(self, img, mask=None, crop_params=None): """ @@ -1100,6 +1120,7 @@ def __call__(self, img, mask=None, crop_params=None): np.ndarray or tuple: Cropped image, or tuple of (image, mask) if mask provided """ if crop_params is None: + draw = self.rng if self.rng is not None else np.random height, width = img.shape[:2] h = max(1, int(height * self.crop_scale[0])) w = max(1, int(width * self.crop_scale[1])) @@ -1113,8 +1134,8 @@ def __call__(self, img, mask=None, crop_params=None): if len(xs) == 0 or len(ys) == 0: # No foreground, fall back to random crop - i = np.random.randint(0, height - h + 1) if height > h else 0 - j = np.random.randint(0, width - w + 1) if width > w else 0 + i = draw.randint(0, height - h + 1) if height > h else 0 + j = draw.randint(0, width - w + 1) if width > w else 0 else: x0, x1 = xs.min(), xs.max() y0, y1 = ys.min(), ys.max() @@ -1126,11 +1147,11 @@ def __call__(self, img, mask=None, crop_params=None): i = max(0, y0 - (h // 2)) j = max(0, x0 - (w // 2)) else: - i = np.random.randint(min_i, max_i + 1) if max_i >= min_i else 0 - j = np.random.randint(min_j, max_j + 1) if max_j >= min_j else 0 + i = draw.randint(min_i, max_i + 1) if max_i >= min_i else 0 + j = draw.randint(min_j, max_j + 1) if max_j >= min_j else 0 else: - i = np.random.randint(0, height - h + 1) if height > h else 0 - j = np.random.randint(0, width - w + 1) if width > w else 0 + i = draw.randint(0, height - h + 1) if height > h else 0 + j = draw.randint(0, width - w + 1) if width > w else 0 else: i, j, h, w = crop_params @@ -1216,12 +1237,14 @@ def __call__(self, img, mask=None): class RandomHorizontalFlipWithParams: - def __init__(self, p=0.5): + def __init__(self, p=0.5, rng=None): """ Args: p (float): Probability of flipping the image + rng: np.random.RandomState (or None for the global module). """ self.p = p + self.rng = rng self.params = {} def __call__(self, img, mask=None, flip=None): @@ -1237,7 +1260,8 @@ def __call__(self, img, mask=None, flip=None): if flip is not None: self.params = {"flip": flip} elif not hasattr(self, "params") or len(self.params) == 0: - flip = np.random.random() < self.p + draw = self.rng if self.rng is not None else np.random + flip = draw.random_sample() < self.p self.params = {"flip": flip} if self.params.get("flip", False): @@ -1251,12 +1275,14 @@ def __call__(self, img, mask=None, flip=None): class RandomVerticalFlipWithParams: - def __init__(self, p=0.5): + def __init__(self, p=0.5, rng=None): """ Args: p (float): Probability of flipping the image + rng: np.random.RandomState (or None for the global module). """ self.p = p + self.rng = rng self.params = {} def __call__(self, img, mask=None, flip=None): @@ -1274,7 +1300,8 @@ def __call__(self, img, mask=None, flip=None): if flip is not None: self.params = {"flip": flip} elif not hasattr(self, "params") or len(self.params) == 0: - flip = np.random.random() < self.p + draw = self.rng if self.rng is not None else np.random + flip = draw.random_sample() < self.p self.params = {"flip": flip} if self.params.get("flip", False): @@ -1289,7 +1316,7 @@ def __call__(self, img, mask=None, flip=None): class RandomRotationWithParams: """Randomly rotate an image.""" - def __init__(self, degrees, p=0.5, reshape=False, mode="reflect", order=2): + def __init__(self, degrees, p=0.5, reshape=False, mode="reflect", order=2, rng=None, pyrng=None): """ Initialize random rotation transform. @@ -1299,6 +1326,8 @@ def __init__(self, degrees, p=0.5, reshape=False, mode="reflect", order=2): reshape (bool): If True, expands output image to fit rotated image mode (str): How to fill the border ('reflect', 'constant', etc) order (int): Interpolation order (0-5) + rng: np.random.RandomState (or None for the global module). + pyrng: random.Random (or None for the global module). """ if isinstance(degrees, (tuple, list)): self.degrees = degrees @@ -1309,6 +1338,8 @@ def __init__(self, degrees, p=0.5, reshape=False, mode="reflect", order=2): self.reshape = reshape self.mode = mode self.order = order + self.rng = rng + self.pyrng = pyrng def __call__(self, img, mask, rotate=None, angle=None, order=None, **kwargs): """ @@ -1325,8 +1356,10 @@ def __call__(self, img, mask, rotate=None, angle=None, order=None, **kwargs): Returns: np.ndarray or tuple: Rotated image, or tuple of (image, mask) if mask provided """ + draw = self.rng if self.rng is not None else np.random + pydraw = self.pyrng if self.pyrng is not None else random if rotate is None: - rotate = np.random.random() < self.p + rotate = draw.random_sample() < self.p self.params = {"rotate": rotate} if not rotate: @@ -1336,10 +1369,10 @@ def __call__(self, img, mask, rotate=None, angle=None, order=None, **kwargs): order = self.order if order is None else order if isinstance(order, (tuple, list)): - order = random.randint(order[0], order[1]) + order = pydraw.randint(order[0], order[1]) if angle is None: - angle = random.uniform(self.degrees[0], self.degrees[1]) + angle = pydraw.uniform(self.degrees[0], self.degrees[1]) self.params.update({"order": order, "angle": angle}) img = ndimage.rotate( diff --git a/tests/unit/test_column_detection.py b/tests/unit/test_column_detection.py index 6bed619..2eaddaa 100644 --- a/tests/unit/test_column_detection.py +++ b/tests/unit/test_column_detection.py @@ -5,8 +5,6 @@ """ import pytest -from unittest.mock import Mock, patch -import pandas as pd from src.gasbench.dataset.config import BenchmarkDatasetConfig diff --git a/tests/unit/test_config_loading.py b/tests/unit/test_config_loading.py index 6662d7f..18da93a 100644 --- a/tests/unit/test_config_loading.py +++ b/tests/unit/test_config_loading.py @@ -7,7 +7,6 @@ import pytest from src.gasbench.dataset.config import ( load_benchmark_datasets_from_yaml, - BenchmarkDatasetConfig, ) diff --git a/tests/unit/test_transforms_determinism.py b/tests/unit/test_transforms_determinism.py new file mode 100644 index 0000000..eb23d1d --- /dev/null +++ b/tests/unit/test_transforms_determinism.py @@ -0,0 +1,63 @@ +"""Determinism guarantees for the augmentation pipeline. + +The prefetch pipelines run augmentation in parallel worker threads. These tests +pin the two properties that must hold for reproducible scoring: + 1. A given (sample, seed) produces the same output regardless of concurrency. + 2. Different seeds produce different augmentations (the seed is actually used). +""" + +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pytest + +from src.gasbench.processing.transforms import ( + apply_random_augmentations, + apply_robustness_augmentations, +) + +TARGET = (224, 224) + + +def _img(seed=0, h=260, w=300): + return (np.random.RandomState(seed).rand(h, w, 3) * 255).astype(np.uint8) + + +@pytest.mark.parametrize("level", [0, 1, 2, 3]) +def test_same_seed_same_output(level): + img = _img() + a, *_ = apply_random_augmentations(img.copy(), TARGET, level=level, crop_prob=0.5, seed=42) + b, *_ = apply_random_augmentations(img.copy(), TARGET, level=level, crop_prob=0.5, seed=42) + assert np.array_equal(a, b) + + +def test_different_seed_different_output(): + img = _img() + a, *_ = apply_random_augmentations(img.copy(), TARGET, level=3, crop_prob=0.5, seed=1) + b, *_ = apply_random_augmentations(img.copy(), TARGET, level=3, crop_prob=0.5, seed=2) + assert not np.array_equal(a, b) + + +def test_deterministic_under_thread_pool(): + """The core guarantee: seeding is per-call, not global, so concurrent workers + cannot clobber each other's RNG state.""" + samples = [(_img(seed=i), 1000 + i) for i in range(24)] + + def one(item): + img, seed = item + out, *_ = apply_random_augmentations(img.copy(), TARGET, level=3, crop_prob=0.5, seed=seed) + return out + + serial = [one(s) for s in samples] + with ThreadPoolExecutor(max_workers=8) as ex: + concurrent = list(ex.map(one, samples)) + + for i, (s, c) in enumerate(zip(serial, concurrent)): + assert np.array_equal(s, c), f"sample {i} differed between serial and concurrent runs" + + +def test_robustness_pass_is_deterministic(): + img = _img() + a, *_ = apply_robustness_augmentations(img.copy(), TARGET, seed=42) + b, *_ = apply_robustness_augmentations(img.copy(), TARGET, seed=42) + assert np.array_equal(a, b)