Spaces:
Sleeping
Sleeping
Download backend/media_analysis/figure_analysis.py from EnvironmentalAI/WaterScopeAI-Search: direct link, hf CLI and curl.
- Browser
- Download file 10.2 kB
-
https://huggingface.co/spaces/EnvironmentalAI/WaterScopeAI-Search/resolve/main/backend/media_analysis/figure_analysis.py
- Command line
-
hf download hf://spaces/EnvironmentalAI/WaterScopeAI-Search/backend/media_analysis/figure_analysis.py
-
curl -L -o figure_analysis.py https://huggingface.co/spaces/EnvironmentalAI/WaterScopeAI-Search/resolve/main/backend/media_analysis/figure_analysis.py
10.2 kB
| """Resolve selected scientific paper media using the EXISTING paper_figures module. | |
| This adapter is isolated from existing chat and routing. Production chat integration | |
| should happen only after verifying IDs, manifest records and endpoint behavior. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import re | |
| from typing import Any, Callable | |
| class MediaSelectionError(ValueError): | |
| pass | |
| def _canonical_label(label: str) -> tuple[str, int] | None: | |
| # Extracted manifest labels use figure_03 / table_04, while users ask | |
| # for Figure 3 / Fig. 3 / Table 4. Normalize both without touching manifests. | |
| match = re.search(r"(?<!\w)(fig(?:ure)?|table)[\s_.-]*(\d+)(?!\d)", label or "", re.I) | |
| if not match: | |
| return None | |
| kind = "table" if match.group(1).lower() == "table" else "figure" | |
| return kind, int(match.group(2)) | |
| def _requested_labels(question: str) -> list[tuple[str, int]]: | |
| found = [] | |
| for match in re.finditer(r"\b(fig(?:ure)?\.?|table)\s*(\d+)\b", question, re.I): | |
| label = ("table" if match.group(1).lower() == "table" else "figure", int(match.group(2))) | |
| if label not in found: | |
| found.append(label) | |
| return found | |
| def prepare_media_evidence( | |
| question: str, | |
| selected_papers: list[dict[str, Any]], | |
| *, | |
| figures_lookup: Callable[..., dict | None] | None = None, | |
| asset_resolver: Callable[[str, str], Any] | None = None, | |
| max_images: int = 4, | |
| media_image: str | None = None, | |
| ) -> list[dict]: | |
| """Explicit label-matching POC; never retrieves unselected papers. | |
| selected_papers must contain paper_id and at least one of paper_key/txt_file/ | |
| pdf_file/metadata_row_index. Optional paper_title, doi, summary are carried | |
| without performing new retrieval. Ambiguous figure labels across selected | |
| papers are deliberately rejected (identify the paper or pass one paper). | |
| """ | |
| if figures_lookup is None or asset_resolver is None: | |
| from metadata_search.paper_figures import figures_payload_for_paper, resolve_figure_asset | |
| figures_lookup = figures_lookup or figures_payload_for_paper | |
| asset_resolver = asset_resolver or resolve_figure_asset | |
| if not selected_papers: | |
| raise MediaSelectionError("Select at least one paper") | |
| if max_images < 1 or max_images > 4: | |
| raise ValueError("max_images must be 1-4 for the POC") | |
| labels = _requested_labels(question) | |
| if not labels: | |
| raise MediaSelectionError("POC requires an explicit Figure N or Table N reference") | |
| if len(selected_papers) > 4: | |
| raise MediaSelectionError("POC supports up to four selected papers") | |
| matched = [] | |
| for paper in selected_papers: | |
| paper_id = str(paper.get("paper_id") or paper.get("document_id") or "").strip() | |
| if not paper_id: | |
| raise MediaSelectionError("Each selected paper needs a stable paper_id/document_id") | |
| lookup_args = {key: paper[key] for key in ("txt_file", "pdf_file", "paper_key", "metadata_row_index") | |
| if paper.get(key) is not None and paper.get(key) != ""} | |
| if not lookup_args: | |
| raise MediaSelectionError(f"Paper {paper_id} missing a media lookup locator") | |
| payload = figures_lookup(**lookup_args) | |
| if not payload: | |
| continue | |
| folder = payload.get("paper_key") | |
| if not folder: | |
| continue | |
| for entry in payload.get("items") or []: | |
| actual = _canonical_label(entry.get("label", "")) | |
| if actual not in labels: | |
| continue | |
| # Reject duplicated figure labels inside one paper rather than guessing. | |
| image = entry.get("image") | |
| if media_image is not None and image != media_image: | |
| continue | |
| if not image: | |
| continue | |
| image_path = asset_resolver(folder, image) | |
| if image_path is None: | |
| raise MediaSelectionError(f"Image missing for paper {paper_id}, {entry.get('label')}") | |
| matched.append({ | |
| "paper_id": paper_id, | |
| "paper_title": paper.get("paper_title") or paper.get("title", ""), | |
| "doi": paper.get("doi", ""), | |
| "summary": paper.get("summary", ""), | |
| "paper_key": folder, | |
| "label": entry.get("label", ""), | |
| "panel": "", # Full crop first; panel-specific cropping is not available. | |
| "page": entry.get("page"), | |
| "caption": entry.get("caption", ""), | |
| "image_url": entry.get("image_url", ""), | |
| "image_path": str(image_path), | |
| "kind": entry.get("kind", "figure"), | |
| }) | |
| if not matched: | |
| raise MediaSelectionError("Requested image not found for this paper and figure/table label" if media_image else "Requested media not found in the selected papers") | |
| # Do not conflate 'Figure 3' across multiple papers without a paper qualifier. | |
| if len(selected_papers) > 1 and len(labels) == 1 and len(matched) > 1: | |
| # Explicit cross-paper comparisons are allowed; otherwise require disambiguation. | |
| if not re.search(r"\b(compare|across|between|both|all|different studies|selected papers)\b", question, re.I): | |
| raise MediaSelectionError("Figure/table label occurs in several selected papers; specify paper or ask an explicit cross-paper comparison") | |
| if len(matched) > max_images: | |
| raise MediaSelectionError(f"Matched {len(matched)} images; POC limit is {max_images}. Narrow the request") | |
| keys = [(x["paper_id"], _canonical_label(x["label"])) for x in matched] | |
| if len(set(keys)) != len(keys): | |
| raise MediaSelectionError("Multiple images share this figure/table number. Click Analyze & CSV under the intended image in the gallery.") | |
| return matched | |
| def analyze_media_question(question: str, selected_papers: list[dict], *, enabled: bool | None = None, | |
| figures_lookup=None, asset_resolver=None, vision_client=None, provider: str = "azure", model: str | None = None, max_output_tokens: int | None = None, | |
| evidence_mode: str = "vision", export_csv: bool = False, analysis_mode: str = "auto", media_image: str | None = None) -> dict: | |
| """Explicit opt-in entry point. Never silently substitutes a text summary answer.""" | |
| if enabled is None: | |
| enabled = os.getenv("WATERSCOPE_ENABLE_MEDIA_ANALYSIS", "0").lower() in ("1", "true", "yes") | |
| if not enabled: | |
| return {"status": "disabled", "answer": None, "vision_inspected": False} | |
| items = prepare_media_evidence(question, selected_papers, | |
| figures_lookup=figures_lookup, asset_resolver=asset_resolver, media_image=media_image) | |
| if vision_client is None: | |
| from .vision_client import analyze_images | |
| vision_client = analyze_images | |
| if analysis_mode not in ("auto", "extract_all", "interpret"): | |
| raise ValueError("analysis_mode must be auto, extract_all, or interpret") | |
| if analysis_mode == "extract_all": | |
| question = question + ("\n[Task: extract all legitimately readable values from the selected media; " | |
| "do not invent chart points or force numerical estimates where unreadable.]") | |
| elif analysis_mode == "interpret": | |
| question = question + "\n[Task: prioritize interpretation; extract only data relevant to the question.]" | |
| if evidence_mode not in ("vision", "json", "combined"): | |
| raise ValueError("evidence_mode must be vision, json or combined") | |
| # Figure images usually have no positional JSON; combined mode automatically | |
| # uses the image and caption without requiring a fictitious figure JSON. | |
| if evidence_mode != "vision": | |
| from .table_data import load_table_evidence | |
| if asset_resolver is None: | |
| from metadata_search.paper_figures import resolve_figure_asset | |
| asset_resolver = resolve_figure_asset | |
| for item in items: | |
| if item.get("kind") == "table": | |
| item["table_data"] = load_table_evidence(item["paper_key"], item["label"], asset_resolver) | |
| if evidence_mode == "json" and any(not x.get("table_data") for x in items): | |
| raise MediaSelectionError("JSON-only mode requires extracted table JSON for every requested item") | |
| result = vision_client(question, items, provider=provider, model=model, | |
| max_output_tokens=max_output_tokens, evidence_mode=evidence_mode, | |
| structured_output=export_csv) | |
| if export_csv: | |
| from .data_export import decode_structured_answer, csv_attachment, normalize_records | |
| if result.get("response_status") == "incomplete" or result.get("incomplete_reason"): | |
| answer, records, warning = "", [], "Model output incomplete; CSV was not generated" | |
| else: | |
| answer, records, warning = decode_structured_answer(result.get("answer", "")) | |
| result["answer"] = answer if not warning else ( | |
| "Structured extraction did not complete; no downloadable CSV was generated. " | |
| "Please retry or narrow the request." | |
| ) | |
| records, normalization_warnings = normalize_records(records) | |
| result["structured_data"] = records | |
| result["normalization_warnings"] = normalization_warnings | |
| sources = result.get("sources") or [] | |
| if len(sources) > 1 and records: | |
| names = {key.lower() for rec in records for key in rec} | |
| if not any(name in names for name in ("source", "source_id", "paper_id", "paper_title", "source_media_label", "media_label", "source_label")): | |
| warning = warning or "Multiple sources require a source identifier column; CSV not generated" | |
| records = [] | |
| result["structured_data"] = [] | |
| result["csv_attachment"] = csv_attachment(records, sources) | |
| if warning: | |
| result["export_warning"] = warning | |
| return {"status": "ok", "vision_inspected": evidence_mode != "json", "evidence_mode": evidence_mode, "analysis_mode": analysis_mode, **result} | |