File size: 6,265 Bytes
15d68eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """
Search HuggingFace Hub for pre-curated Indian heritage art datasets.
HuggingFace hosts user-contributed datasets that are perfect for LoRA training.
This script searches for relevant datasets and lets you download them in one shot.
Usage:
python scripts/search_hf_datasets.py
python scripts/search_hf_datasets.py --download madhubani-art
"""
from __future__ import annotations
import argparse
import logging
from pathlib import Path
log = logging.getLogger(__name__)
# Known datasets worth checking (will verify and update at runtime)
KNOWN_DATASETS = [
# Indian art / painting datasets
"rishiMadhuri/madhubani-art",
"MMadhubani/madhubani-painting",
"ravi105/madhubani-art",
"nerdlab/indian-art",
"shubhamorhan/indian-painting-styles",
"Sakshi1307/indian-art-forms",
"DrishtiSharma/indian-heritage-art",
"rcss/indian-art-classification",
# General art datasets that include Indian art
"huggan/few-shot-art",
"huggan/wikiart",
"Artificio/WikiArt",
# Bharat-specific
"Bharat/desi-art",
"kabita-chatterjee/indian-art-forms",
]
# Search queries to run on the HF Hub API
SEARCH_QUERIES = [
"madhubani",
"warli",
"pattachitra",
"mughal painting",
"tanjore painting",
"indian art",
"indian heritage",
"indian painting",
"folk art india",
"traditional art india",
]
def search_hf_datasets(query: str, limit: int = 10) -> list:
"""Search HuggingFace Hub for datasets matching a query."""
try:
from huggingface_hub import HfApi
api = HfApi()
results = list(api.list_datasets(search=query, limit=limit))
return [
{
"id": d.id,
"downloads": d.downloads,
"likes": d.likes,
"tags": d.tags,
}
for d in results
]
except Exception as exc:
log.error(f"HF search failed for '{query}': {exc}")
return []
def check_dataset_contents(dataset_id: str) -> dict:
"""Check if a dataset has images and what splits/columns it has."""
try:
from datasets import load_dataset_builder
builder = load_dataset_builder(dataset_id)
info = {
"id": dataset_id,
"description": (builder.info.description or "")[:200],
"features": str(builder.info.features)[:500],
"splits": list(builder.info.splits.keys()) if builder.info.splits else [],
"size": builder.info.dataset_size or "unknown",
}
return info
except Exception as exc:
return {"id": dataset_id, "error": str(exc)[:200]}
def download_dataset(dataset_id: str, out_dir: Path, max_images: int = 50) -> int:
"""Download a dataset and save images to out_dir. Returns count saved."""
out_dir.mkdir(parents=True, exist_ok=True)
try:
from datasets import load_dataset
ds = load_dataset(dataset_id, split="train")
saved = 0
for i, item in enumerate(ds):
if saved >= max_images:
break
# Find the image field
img = None
for key in ("image", "img", "painting", "art", "picture"):
if key in item:
img = item[key]
break
if img is None:
# Try first PIL Image field
for k, v in item.items():
if hasattr(v, "save"):
img = v
break
if img is None:
continue
try:
img = img.convert("RGB")
out_path = out_dir / f"{dataset_id.replace('/', '_')}_{saved:03d}.jpg"
img.save(out_path, "JPEG", quality=95)
saved += 1
except Exception:
continue
return saved
except Exception as exc:
log.error(f"Download failed for {dataset_id}: {exc}")
return 0
def main():
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
p = argparse.ArgumentParser()
p.add_argument("--query", help="Search for datasets matching this query")
p.add_argument("--download", help="Download this dataset ID")
p.add_argument("--style", help="Save downloads to assets/datasets/raw/<style>/")
p.add_argument("--max", type=int, default=50, help="Max images to download")
p.add_argument("--check", help="Check contents of this dataset ID")
args = p.parse_args()
if args.check:
info = check_dataset_contents(args.check)
print(f"\n=== {args.check} ===")
for k, v in info.items():
print(f" {k}: {v}")
return
if args.download:
out_dir = Path(f"assets/datasets/raw/{args.style}") if args.style else Path("outputs/hf_download")
n = download_dataset(args.download, out_dir, max_images=args.max)
print(f"Downloaded {n} images to {out_dir}")
return
# Default: search across all queries
print("=" * 60)
print("HuggingFace Heritage Art Dataset Search")
print("=" * 60)
all_results = {}
for q in SEARCH_QUERIES:
print(f"\n--- Search: '{q}' ---")
results = search_hf_datasets(q, limit=5)
for r in results:
print(f" {r['id']:50s} | ↓{r['downloads']:>6} | ♥{r['likes']:>3} | tags: {[t for t in r.get('tags', []) if 'task' in t][:2]}")
all_results[r["id"]] = r
# Also check known datasets
print("\n--- Checking known datasets ---")
for ds_id in KNOWN_DATASETS:
info = check_dataset_contents(ds_id)
if "error" in info:
print(f" ✗ {ds_id}: {info['error'][:80]}")
else:
print(f" ✓ {ds_id}")
print(f" splits: {info['splits']}")
print(f" features: {info['features'][:200]}")
print(f"\n=== Found {len(all_results)} candidate datasets ===")
print("To check a dataset's contents:")
print(f" python scripts/search_hf_datasets.py --check <dataset-id>")
print("\nTo download a dataset:")
print(f" python scripts/search_hf_datasets.py --download <dataset-id> --style madhubani --max 50")
if __name__ == "__main__":
main()
|