scripts/search_hf_marine_datasets.py

"""Search Hugging Face for marine feature datasets and write normalized manifests.

The script does not download full datasets. It inspects dataset repository
metadata and file lists, then writes project-compatible manifests using hf://
paths so large public datasets can be reviewed before any costly download.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable

from huggingface_hub import HfApi


IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".jp2", ".bmp", ".webp"}
METADATA_EXTS = {".csv", ".json", ".jsonl", ".parquet", ".txt"}
ARCHIVE_EXTS = {".zip", ".tar", ".gz", ".tgz", ".7z"}
TRACKED_EXTS = IMAGE_EXTS | METADATA_EXTS | ARCHIVE_EXTS
MASK_HINTS = ("mask", "label", "labels", "gt", "annotation", "annotations", "seg", "target")
TRAIN_SPLITS = ("train", "training", "val", "valid", "validation", "test")

QUERIES = [
    "seaweed",
    "green tide",
    "red tide",
    "sargassum",
    "aquaculture",
    "ship",
    "oil spill",
    "sea ice",
    "marine",
    "ocean",
    "remote sensing",
    "sar ship",
    "satellite imagery",
]

ELEMENT_KEYWORDS = {
    "green_tide": ("green_tide", "greentide", "green tide", "enteromorpha", "seaweed"),
    "red_tide": ("red_tide", "redtide", "red tide", "harmful algal", "hab"),
    "golden_tide": ("golden_tide", "goldentide", "sargassum", "sarg"),
    "aquaculture": ("aquaculture", "oyster", "raft", "fish cage", "fish-cage"),
    "ship": ("ship", "sar ship"),
    "oil_spill": ("oilspill", "oil_spill", "oil spill", "oil-spill"),
    "sea_ice": ("seaice", "sea_ice", "sea ice", "arctic ice"),
}

SHIP_CONTEXT_KEYWORDS = (
    "ship",
    "sar",
    "satellite",
    "remote sensing",
    "marine",
    "ocean",
    "sentinel",
    "ais",
)

EXCLUDE_KEYWORDS = (
    "medical",
    "retinal",
    "retina",
    "miccai",
    "flare",
    "drive digital retinal",
    "godot",
    "shipping law",
    "shipping orders",
    "textual inversion",
    "kantaicollection",
    "waifu",
    "anime",
    "cancer",
    "community health",
    "plankton",
    "mammal",
    "legal",
)

SATELLITE_PATTERN = re.compile(r"\b(GF\d+|HY\d+|Sentinel-?1|Sentinel-?2|Landsat-?\d*|SAR)\b", re.IGNORECASE)
PATCH_SIZE_PATTERN = re.compile(r"(?:^|[_/\-])(?:size)?(128|256|512|1024)(?:[_/\-]|$)")


@dataclass
class HfAssetRecord:
    asset_id: str
    repo_id: str
    path: str
    hf_path: str
    filename: str
    suffix: str
    role: str
    element: str
    satellite: str | None
    sensor: str | None
    patch_size: int | None
    source_project: str
    source_dataset: str
    size_bytes: int | None
    downloads: int | None
    likes: int | None
    license: str | None
    tags: list[str]
    discovered_by: list[str]
    quality_flags: list[str]


@dataclass
class HfSampleRecord:
    sample_id: str
    element: str
    task_type: str
    image_path: str
    mask_path: str | None
    label_encoding: dict[str, str] | None
    satellite: str | None
    sensor: str | None
    resolution_m: float | None
    patch_size: int | None
    bands: list[str] | None
    band_count: int | None
    dtype: str | None
    fusion: dict
    acquired_at: str | None
    source_project: str
    source_dataset: str
    split: str | None
    quality_flags: list[str]
    notes: str


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-root", required=True)
    parser.add_argument("--author", default="cuibinge", help="HF namespace to always include.")
    parser.add_argument("--limit-per-query", type=int, default=20)
    parser.add_argument("--max-repos", type=int, default=80)
    parser.add_argument("--max-files-per-repo", type=int, default=5000)
    parser.add_argument("--include-low-confidence", action="store_true")
    return parser.parse_args()


def normalize_text(value: str) -> str:
    return value.replace("_", " ").replace("-", " ").replace("/", " ").lower()


def infer_element(*parts: str) -> str:
    text = normalize_text(" ".join(part for part in parts if part))
    if "vessel" in text and all(keyword not in text for keyword in SHIP_CONTEXT_KEYWORDS):
        return "unknown"
    for element, keywords in ELEMENT_KEYWORDS.items():
        if any(keyword in text for keyword in keywords):
            return element
    if "vessel" in text and any(keyword in text for keyword in SHIP_CONTEXT_KEYWORDS):
        return "ship"
    return "unknown"


def infer_role(path: str) -> str:
    lower = path.lower()
    stem = Path(path).stem.lower()
    suffix = Path(path).suffix.lower()
    if suffix in ARCHIVE_EXTS:
        return "archive"
    if suffix in METADATA_EXTS:
        if any(hint in lower or hint in stem for hint in MASK_HINTS) or stem in {"train", "val", "test"}:
            return "annotation_table"
        return "metadata"
    if any(hint in lower or hint in stem for hint in MASK_HINTS):
        return "mask"
    return "image"


def infer_satellite(*parts: str) -> str | None:
    match = SATELLITE_PATTERN.search(" ".join(parts))
    return match.group(1).upper().replace("-", "") if match else None


def infer_sensor(*parts: str) -> str | None:
    upper = " ".join(parts).upper()
    for sensor in ("PMS", "MUX", "MSS", "PAN", "WFV", "SAR", "MSI", "OLI"):
        if sensor in upper:
            return sensor
    return None


def infer_patch_size(path: str) -> int | None:
    match = PATCH_SIZE_PATTERN.search(path)
    return int(match.group(1)) if match else None


def infer_split(path: str) -> str | None:
    parts = {part.lower() for part in Path(path).parts}
    for split in TRAIN_SPLITS:
        if split in parts:
            return "val" if split in {"valid", "validation"} else "train" if split == "training" else split
    return None


def infer_license(tags: list[str]) -> str | None:
    for tag in tags:
        if tag.startswith("license:"):
            return tag.split(":", 1)[1]
    return None


def infer_fusion(repo_id: str, path: str, sensor: str | None) -> dict:
    lower = f"{repo_id}/{path}".lower()
    if "fuse" in lower or "fusion" in lower:
        state = "fused_product"
        method = "unknown_vendor_product"
        persisted = True
    elif sensor in {"SAR", "MSI", "OLI"}:
        state = "none"
        method = "none"
        persisted = False
    else:
        state = "unknown"
        method = "unknown"
        persisted = False
    return {
        "state": state,
        "method": method,
        "sources": [{"role": "hf_dataset_file", "path": f"hf://datasets/{repo_id}/{path}", "resolution_m": None}],
        "target_resolution_m": None,
        "native_multispectral_resolution_m": None,
        "persisted": persisted,
        "reproducible": False,
        "spectral_preservation": "unknown",
        "notes": "Inferred from Hugging Face repository metadata; verify before training.",
    }


def safe_id(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9]+", "_", value).strip("_").lower()[:180]


def dataset_search(api: HfApi, author: str, limit_per_query: int, max_repos: int) -> dict[str, dict]:
    repos: dict[str, dict] = {}

    def add_repo(dataset, reason: str) -> None:
        entry = repos.setdefault(dataset.id, {"dataset": dataset, "reasons": []})
        if reason not in entry["reasons"]:
            entry["reasons"].append(reason)

    for dataset in api.list_datasets(author=author, full=True):
        add_repo(dataset, f"author:{author}")

    for query in QUERIES:
        for dataset in api.list_datasets(search=query, limit=limit_per_query, full=True):
            add_repo(dataset, f"query:{query}")
            if len(repos) >= max_repos:
                return repos
    return repos


def relevant_repo(repo_id: str, tags: list[str], reasons: list[str]) -> bool:
    text = normalize_text(" ".join([repo_id, " ".join(tags), " ".join(reasons)]))
    if any(keyword in text for keyword in EXCLUDE_KEYWORDS):
        return False
    return any(keyword in text for keywords in ELEMENT_KEYWORDS.values() for keyword in keywords) or any(
        token in text for token in ("marine", "ocean", "remote sensing", "satellite", "sar", "coast", "sea land")
    )


def iter_siblings(api: HfApi, repo_id: str, max_files: int) -> Iterable:
    info = api.dataset_info(repo_id=repo_id, files_metadata=True)
    for idx, sibling in enumerate(info.siblings or []):
        if idx >= max_files:
            break
        yield sibling


def write_jsonl(path: Path, rows: Iterable[dict]) -> None:
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def main() -> None:
    args = parse_args()
    output_root = Path(args.output_root)
    manifest_dir = output_root / "manifests"
    report_dir = output_root / "reports"
    manifest_dir.mkdir(parents=True, exist_ok=True)
    report_dir.mkdir(parents=True, exist_ok=True)

    api = HfApi(token=os.environ.get("HF_TOKEN"))
    repos = dataset_search(api, args.author, args.limit_per_query, args.max_repos)

    assets: list[HfAssetRecord] = []
    repo_rows: list[dict] = []

    for repo_id, entry in sorted(repos.items()):
        dataset = entry["dataset"]
        tags = list(getattr(dataset, "tags", []) or [])
        reasons = entry["reasons"]
        if not args.include_low_confidence and not relevant_repo(repo_id, tags, reasons):
            continue

        try:
            siblings = list(iter_siblings(api, repo_id, args.max_files_per_repo))
        except Exception as exc:  # noqa: BLE001 - keep discovery resilient.
            repo_rows.append({"repo_id": repo_id, "status": "error", "error": str(exc), "reasons": ";".join(reasons)})
            continue

        image_count = 0
        asset_count = 0
        for sibling in siblings:
            rel_path = sibling.rfilename
            suffix = Path(rel_path).suffix.lower()
            if suffix not in TRACKED_EXTS:
                continue
            role = infer_role(rel_path)
            if suffix in IMAGE_EXTS:
                image_count += 1
            element = infer_element(repo_id, rel_path, " ".join(tags))
            sensor = infer_sensor(repo_id, rel_path, " ".join(tags))
            flags: list[str] = []
            if element == "unknown":
                flags.append("unknown_element")
            if role in {"mask", "annotation_table"}:
                flags.append("mask_asset")
            assets.append(
                HfAssetRecord(
                    asset_id=safe_id(f"{repo_id}_{rel_path}"),
                    repo_id=repo_id,
                    path=rel_path,
                    hf_path=f"hf://datasets/{repo_id}/{rel_path}",
                    filename=Path(rel_path).name,
                    suffix=suffix,
                    role=role,
                    element=element,
                    satellite=infer_satellite(repo_id, rel_path, " ".join(tags)),
                    sensor=sensor,
                    patch_size=infer_patch_size(rel_path),
                    source_project=repo_id.split("/", 1)[-1],
                    source_dataset=repo_id,
                    size_bytes=getattr(sibling, "size", None),
                    downloads=getattr(dataset, "downloads", None),
                    likes=getattr(dataset, "likes", None),
                    license=infer_license(tags),
                    tags=tags,
                    discovered_by=reasons,
                    quality_flags=flags,
                )
            )
            asset_count += 1

        repo_rows.append(
            {
                "repo_id": repo_id,
                "status": "ok",
                "reasons": ";".join(reasons),
                "downloads": getattr(dataset, "downloads", None),
                "likes": getattr(dataset, "likes", None),
                "license": infer_license(tags),
                "tags": ";".join(tags[:20]),
                "files_scanned": len(siblings),
                "image_assets": image_count,
                "normalized_assets": asset_count,
            }
        )

    masks_by_key: dict[tuple[str, str], HfAssetRecord] = {}
    for asset in assets:
        if asset.role == "mask":
            masks_by_key[(asset.repo_id, Path(asset.path).stem.lower())] = asset

    samples: list[HfSampleRecord] = []
    for asset in assets:
        if asset.role != "image":
            continue
        mask = masks_by_key.get((asset.repo_id, Path(asset.path).stem.lower()))
        flags = list(asset.quality_flags)
        if mask is None:
            flags.append("unpaired_hf_image")
        if asset.element == "unknown":
            flags.append("needs_element_review")
        samples.append(
            HfSampleRecord(
                sample_id=f"hf_{safe_id(asset.repo_id)}_{asset.asset_id}",
                element=asset.element,
                task_type="semantic_segmentation" if mask else "image_asset",
                image_path=asset.hf_path,
                mask_path=mask.hf_path if mask else None,
                label_encoding={"0": "background", "1": asset.element} if mask and asset.element != "unknown" else None,
                satellite=asset.satellite,
                sensor=asset.sensor,
                resolution_m=None,
                patch_size=asset.patch_size,
                bands=None,
                band_count=None,
                dtype=None,
                fusion=infer_fusion(asset.repo_id, asset.path, asset.sensor),
                acquired_at=None,
                source_project=asset.source_project,
                source_dataset=asset.source_dataset,
                split=infer_split(asset.path),
                quality_flags=flags,
                notes="HF-discovered asset; inspect license, labels, georeferencing, and split before training.",
            )
        )

    ready_samples = [
        sample
        for sample in samples
        if sample.element != "unknown" and sample.task_type == "semantic_segmentation" and sample.mask_path
    ]
    review_samples = [sample for sample in samples if sample not in ready_samples]

    write_jsonl(manifest_dir / "hf_assets_raw.jsonl", (asdict(asset) for asset in assets))
    write_jsonl(manifest_dir / "samples.jsonl", (asdict(sample) for sample in samples))
    write_jsonl(manifest_dir / "samples_ready.jsonl", (asdict(sample) for sample in ready_samples))
    write_jsonl(manifest_dir / "samples_review.jsonl", (asdict(sample) for sample in review_samples))

    with (report_dir / "hf_dataset_inventory.csv").open("w", newline="", encoding="utf-8-sig") as f:
        fieldnames = [
            "repo_id",
            "status",
            "reasons",
            "downloads",
            "likes",
            "license",
            "tags",
            "files_scanned",
            "image_assets",
            "normalized_assets",
            "error",
        ]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for row in repo_rows:
            writer.writerow({name: row.get(name, "") for name in fieldnames})

    summary = {
        "created_at": datetime.now(timezone.utc).isoformat(),
        "repo_count": len(repo_rows),
        "asset_count": len(assets),
        "sample_count": len(samples),
        "ready_sample_count": len(ready_samples),
        "review_sample_count": len(review_samples),
        "by_element": {},
        "ready_by_element": {},
        "output_root": str(output_root),
        "notes": [
            "hf:// paths are references; full dataset download is intentionally deferred.",
            "Unknown elements and unpaired images require manual review before training.",
            "No external coastline or land-mask vector is assumed.",
        ],
    }
    for sample in samples:
        summary["by_element"][sample.element] = summary["by_element"].get(sample.element, 0) + 1
    for sample in ready_samples:
        summary["ready_by_element"][sample.element] = summary["ready_by_element"].get(sample.element, 0) + 1
    (report_dir / "hf_dataset_summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")

    md_lines = [
        "# Hugging Face Marine Dataset Discovery",
        "",
        f"- Created: {summary['created_at']}",
        f"- Repositories reviewed: {summary['repo_count']}",
        f"- Image assets indexed: {summary['asset_count']}",
        f"- Standardized samples written: {summary['sample_count']}",
        f"- Training-ready samples: {summary['ready_sample_count']}",
        f"- Review samples: {summary['review_sample_count']}",
        "",
        "## Samples By Element",
        "",
    ]
    for element, count in sorted(summary["by_element"].items()):
        md_lines.append(f"- `{element}`: {count}")
    md_lines.extend(["", "## Training-Ready Samples By Element", ""])
    for element, count in sorted(summary["ready_by_element"].items()):
        md_lines.append(f"- `{element}`: {count}")
    md_lines.extend(
        [
            "",
            "## Important Notes",
            "",
            "- Manifests use `hf://datasets/<repo>/<path>` references and do not imply files were downloaded.",
            "- Licenses and label semantics must be checked before a dataset is used for training.",
            "- Unknown or unpaired assets are kept for review, not treated as training-ready negatives.",
        ]
    )
    (report_dir / "hf_dataset_discovery.md").write_text("\n".join(md_lines) + "\n", encoding="utf-8")
    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()