worldmodel/research/fetch_crowdroom_papers.py

"""CrowdRoom 相关 arXiv 论文抓取脚本。

复用 research/search_info.py 中的 search_arxiv 与 build_opener，
定义 CrowdRoom 项目专属的 10 组查询主题。

用法：
    HTTPS_PROXY=http://127.0.0.1:6984 HTTP_PROXY=http://127.0.0.1:6984 \
        python3 research/fetch_crowdroom_papers.py \
        --proxy http://127.0.0.1:6984 \
        --max-results 20 --delay 5.0 \
        --out research/crowdroom_papers_raw.json

参数：
    --proxy        代理（默认走环境变量 / 直连）
    --max-results  每主题最大论文数（默认 20）
    --delay        主题间延迟秒数（arXiv 政策 ≥3，建议 5）
    --topics       只跑指定主题
    --out          输出 JSON
    --max-days     仅保留 published 在 N 天内的论文（默认 365；0 = 不过滤）
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone, timedelta

# 复用 search_info.py
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from search_info import build_opener  # noqa: E402

# 自实现 search_arxiv，带更激进的退避（针对 arxiv 429 限速窗口）
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET


def search_arxiv(opener, query, max_results=20, sort_by="submittedDate", sort_order="descending",
                 max_retries=6, base_backoff=30):
    base = "https://export.arxiv.org/api/query"
    params = urllib.parse.urlencode({
        "search_query": query, "start": 0, "max_results": max_results,
        "sortBy": sort_by, "sortOrder": sort_order,
    })
    url = f"{base}?{params}"
    req = urllib.request.Request(url, headers={
        "User-Agent": "crowdroom-research-bot/0.1 (mailto:research@example.com)"
    })

    data = ""
    for attempt in range(max_retries):
        try:
            with opener.open(req, timeout=45) as resp:
                data = resp.read().decode("utf-8")
            break
        except urllib.error.HTTPError as e:
            if e.code == 429:
                wait = base_backoff * (attempt + 1)
                print(f"[arxiv][{query[:40]}...] 429 限速，等待 {wait}s 重试 ({attempt+1}/{max_retries})", file=sys.stderr)
                time.sleep(wait)
                continue
            print(f"[arxiv][{query[:40]}...] HTTPError {e.code}", file=sys.stderr)
            return []
        except (urllib.error.URLError, TimeoutError) as e:
            wait = 15 * (attempt + 1)
            print(f"[arxiv][{query[:40]}...] timeout/url error: {e}; 等待 {wait}s 重试", file=sys.stderr)
            time.sleep(wait)
            continue
    else:
        print(f"[arxiv][{query[:40]}...] 持续失败，放弃", file=sys.stderr)
        return []

    try:
        root = ET.fromstring(data)
    except ET.ParseError as e:
        print(f"[arxiv][{query[:40]}...] XML parse error: {e}", file=sys.stderr)
        return []

    ns = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}
    papers = []
    for entry in root.findall("atom:entry", ns):
        def _text(elem_path):
            elem = entry.find(elem_path, ns)
            return (elem.text or "").replace("\n", " ").strip() if elem is not None else ""

        arxiv_id_url = _text("atom:id")
        arxiv_id = arxiv_id_url.rsplit("/", 1)[-1] if arxiv_id_url else ""
        arxiv_id_clean = arxiv_id.split("v")[0] if arxiv_id else ""

        authors = []
        for a in entry.findall("atom:author", ns):
            name_elem = a.find("atom:name", ns)
            if name_elem is not None and name_elem.text:
                authors.append(name_elem.text.strip())

        categories = [c.attrib.get("term", "") for c in entry.findall("atom:category", ns)]
        summary = _text("atom:summary")

        # 提取 arxiv 专属字段
        primary_cat = ""
        pc = entry.find("arxiv:primary_category", ns)
        if pc is not None:
            primary_cat = pc.attrib.get("term", "")

        # 提取 doi / journal_ref / comment 链接
        pdf_url = ""
        for link in entry.findall("atom:link", ns):
            if link.attrib.get("title") == "pdf":
                pdf_url = link.attrib.get("href", "")

        papers.append({
            "arxiv_id": arxiv_id_clean,
            "url": f"https://arxiv.org/abs/{arxiv_id_clean}" if arxiv_id_clean else arxiv_id_url,
            "pdf_url": pdf_url,
            "title": _text("atom:title"),
            "published": _text("atom:published"),
            "updated": _text("atom:updated"),
            "authors": authors,
            "categories": categories,
            "primary_category": primary_cat,
            "summary": summary,
        })
    return papers


# ──────────────────────────────────────────────────────────────────────
# CrowdRoom 10 组检索主题
# ──────────────────────────────────────────────────────────────────────
TOPICS: dict[str, dict] = {
    "roomplan_arkit": {
        "zh": "1. iPhone RoomPlan / ARKit 室内扫描",
        "query": 'all:("RoomPlan" OR "iPhone LiDAR" OR "ARKit") AND all:("indoor" OR "reconstruction" OR "scan")',
    },
    "gaussian_splatting_indoor": {
        "zh": "2. 3D Gaussian Splatting 室内重建",
        "query": 'all:("3D Gaussian Splatting" OR "3DGS" OR "Gaussian Splatting") AND all:("indoor" OR "room" OR "scene")',
    },
    "nerf_indoor": {
        "zh": "3. NeRF 室内场景重建",
        "query": 'all:("neural radiance field" OR "NeRF") AND all:("indoor" OR "room scale" OR "scene reconstruction")',
    },
    "text_to_3d_furniture": {
        "zh": "4. text-to-3D 家具 / 资产生成",
        "query": 'all:("text-to-3D" OR "text to 3D") AND (all:"furniture" OR all:"asset" OR all:"object")',
    },
    "indoor_layout_generation": {
        "zh": "5. 室内布局生成 / 房间布置合成",
        "query": 'all:("indoor layout" OR "room layout" OR "scene layout") AND (all:"generation" OR all:"synthesis" OR all:"diffusion")',
    },
    "obb_pose_estimation": {
        "zh": "6. 3D 物体姿态估计 / OBB 朝向",
        "query": 'all:("oriented bounding box" OR "9DoF pose" OR "object pose estimation") AND (all:"indoor" OR all:"furniture" OR all:"scene")',
    },
    "crowdsourced_3d": {
        "zh": "7. 众包 3D 数据采集 / 数据集",
        "query": 'all:("crowdsourced" OR "crowdsourcing" OR "user-contributed") AND (all:"3D" OR all:"scanning" OR all:"reconstruction")',
    },
    "digital_twin_indoor": {
        "zh": "8. 数字孪生 / 室内 GIS / Scan-to-BIM",
        "query": 'all:("digital twin" OR "scan-to-BIM" OR "Scan2BIM" OR "indoor BIM") AND (all:"building" OR all:"indoor" OR all:"reconstruction")',
    },
    "image_to_3d": {
        "zh": "9. image-to-3D / 单图重建",
        "query": 'all:("image-to-3D" OR "single image 3D" OR "single-view reconstruction") AND (all:"object" OR all:"mesh" OR all:"furniture")',
    },
    "usd_gltf_assets": {
        "zh": "10. USD / glTF / 3D 资产标准化",
        "query": 'all:("glTF" OR "USD" OR "Universal Scene Description" OR "OpenUSD") AND (all:"3D" OR all:"asset" OR all:"scene")',
    },
}


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="CrowdRoom 相关 arXiv 抓取")
    ap.add_argument("--proxy", default=None)
    ap.add_argument("--max-results", type=int, default=20)
    ap.add_argument("--delay", type=float, default=5.0, help="主题间延迟秒数，arXiv ≥3s")
    ap.add_argument("--topics", nargs="+", default=None)
    ap.add_argument("--out", default="research/crowdroom_papers_raw.json")
    ap.add_argument("--max-days", type=int, default=365, help="仅保留近 N 天提交的论文；0 不过滤")
    ap.add_argument("--sort", choices=["submittedDate", "relevance"], default="submittedDate")
    return ap.parse_args()


def filter_by_date(papers: list[dict], max_days: int) -> list[dict]:
    if max_days <= 0:
        return papers
    cutoff = datetime.now(timezone.utc) - timedelta(days=max_days)
    out = []
    for p in papers:
        pub = p.get("published", "")
        try:
            dt = datetime.fromisoformat(pub.replace("Z", "+00:00"))
        except Exception:
            out.append(p)
            continue
        if dt >= cutoff:
            out.append(p)
    return out


def main() -> int:
    args = parse_args()
    opener = build_opener(args.proxy)
    selected = args.topics or list(TOPICS.keys())

    results: dict = {
        "meta": {
            "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            "proxy": args.proxy or os.environ.get("HTTPS_PROXY") or "direct",
            "max_results_per_topic": args.max_results,
            "max_days": args.max_days,
            "sort": args.sort,
            "topic_count": len(selected),
        },
        "topics_meta": {k: TOPICS[k] for k in selected if k in TOPICS},
        "arxiv": {},
        "stats": {},
    }

    total_raw = 0
    total_after_date = 0

    for i, topic in enumerate(selected, 1):
        if topic not in TOPICS:
            print(f"[skip] 未知主题: {topic}", file=sys.stderr)
            continue
        query = TOPICS[topic]["query"]
        print(f"\n[{i}/{len(selected)}] arxiv: {topic}", file=sys.stderr)
        print(f"        query: {query}", file=sys.stderr)
        papers = search_arxiv(opener, query, max_results=args.max_results, sort_by=args.sort)
        raw_n = len(papers)
        papers = filter_by_date(papers, args.max_days)
        total_raw += raw_n
        total_after_date += len(papers)
        results["arxiv"][topic] = papers
        results["stats"][topic] = {"raw": raw_n, "after_date_filter": len(papers)}
        print(f"        → 抓到 {raw_n} 篇，时间窗内 {len(papers)} 篇", file=sys.stderr)
        if i < len(selected):
            time.sleep(args.delay)

    # 去重统计（全局）
    seen_ids: set[str] = set()
    unique_papers: list[dict] = []
    for topic, plist in results["arxiv"].items():
        for p in plist:
            aid = p.get("arxiv_id", "")
            if not aid or aid in seen_ids:
                continue
            seen_ids.add(aid)
            unique_papers.append({**p, "_topic": topic})

    results["meta"]["total_raw"] = total_raw
    results["meta"]["total_after_date_filter"] = total_after_date
    results["meta"]["total_unique"] = len(unique_papers)
    results["unique_papers"] = unique_papers

    out_path = args.out
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n✅ 完成：raw={total_raw} 时间窗后={total_after_date} 去重={len(unique_papers)}", file=sys.stderr)
    print(f"   → {out_path}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())