raven/utils.py

from __future__ import annotations
import hashlib
import json
import os
import random
from pathlib import Path
from typing import Iterable
import numpy as np
import torch

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def configure_cpu(threads: int = 6, interop_threads: int = 1) -> None:
    threads = max(1, int(threads))
    os.environ.setdefault("OMP_NUM_THREADS", str(threads))
    os.environ.setdefault("MKL_NUM_THREADS", str(threads))
    torch.set_num_threads(threads)
    try:
        torch.set_num_interop_threads(max(1, int(interop_threads)))
    except RuntimeError:
        pass
    torch.backends.mkldnn.enabled = True


def stable_unit(text: str) -> float:
    h = hashlib.sha256(text.encode("utf-8")).digest()
    return int.from_bytes(h[:8], "big") / float(2**64 - 1)


def iter_images(root: str | Path) -> Iterable[Path]:
    root = Path(root)
    if not root.exists():
        return
    for p in root.rglob("*"):
        if p.is_file() and p.suffix.lower() in IMAGE_EXTS:
            yield p


def read_jsonl(path: str | Path) -> list[dict]:
    rows: list[dict] = []
    with Path(path).open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


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


def safe_source_name(source: str) -> str:
    return "".join(c if c.isalnum() or c in "-_." else "_" for c in str(source))[:96]