from __future__ import annotations import os from contextlib import contextmanager from dataclasses import asdict, dataclass, fields, replace from importlib.resources import files from pathlib import Path from typing import Any, Optional from .io import load_image, normalize_inputs, render_pdf from .runtime import resolve_model from ..utils import postprocess_caption DEFAULT_CONFIG = "config.yaml" @dataclass class MolParserConfig: moldet_hf_repo: str = "UniParser/MolDetv2" moldet_modelscope_repo: str = "UniParser/MolDetv2" moldet_image_model_path: str = "" moldet_pdf_model_path: str = "" moldet_image_modelname: str = "moldet_v2_yolo11n_640_general.pt" moldet_pdf_modelname: str = "moldet_v2_yolo11n_960_doc.pt" molparser_hf_repo: str = "UniParser/MolParser-Mobile" molparser_modelscope_repo: str = "UniParser/MolParser-Mobile" molparser_model_path: str = "" cache_dir: str = "" device: str = "auto" detector_conf: float = 0.5 image_imgsz: int = 640 pdf_imgsz: int = 960 moldet_batch_size: int = 8 pdf_dpi: int = 200 max_length: int = 256 molparser_batch_size: int = 32 expand_px: int = 2 padding_px: int = 0 hf_token: str | bool | None = None @classmethod def from_default(cls, **overrides) -> "MolParserConfig": return cls.from_yaml(files("molparser.models").joinpath(DEFAULT_CONFIG), **overrides) @classmethod def from_yaml(cls, path: str | Path, **overrides) -> "MolParserConfig": try: import yaml except ImportError as exc: raise ImportError("YAML config loading requires PyYAML.") from exc with Path(path).open("r", encoding="utf-8") as f: data = yaml.safe_load(f) or {} return cls.from_mapping(data, **overrides) @classmethod def from_mapping(cls, data: dict[str, Any], **overrides) -> "MolParserConfig": flat = { "moldet_hf_repo": data.get("moldet", {}).get("hf_repo", cls.moldet_hf_repo), "moldet_modelscope_repo": data.get("moldet", {}).get("modelscope_repo", cls.moldet_modelscope_repo), "moldet_image_model_path": data.get("moldet", {}).get("image_model_path", ""), "moldet_pdf_model_path": data.get("moldet", {}).get("pdf_model_path", ""), "moldet_image_modelname": data.get("moldet", {}).get("image_modelname", cls.moldet_image_modelname), "moldet_pdf_modelname": data.get("moldet", {}).get("pdf_modelname", cls.moldet_pdf_modelname), "detector_conf": data.get("moldet", {}).get("confidence", cls.detector_conf), "image_imgsz": data.get("moldet", {}).get("image_imgsz", cls.image_imgsz), "pdf_imgsz": data.get("moldet", {}).get("pdf_imgsz", cls.pdf_imgsz), "moldet_batch_size": data.get("moldet", {}).get("batch_size", cls.moldet_batch_size), "expand_px": data.get("moldet", {}).get("expand_px", cls.expand_px), "molparser_hf_repo": data.get("molparser", {}).get("hf_repo", cls.molparser_hf_repo), "molparser_modelscope_repo": data.get("molparser", {}).get("modelscope_repo", ""), "molparser_model_path": data.get("molparser", {}).get("model_path", ""), "max_length": data.get("molparser", {}).get("max_length", cls.max_length), "molparser_batch_size": data.get("molparser", {}).get("batch_size", cls.molparser_batch_size), "padding_px": data.get("molparser", {}).get("padding_px", cls.padding_px), "pdf_dpi": data.get("pdf", {}).get("dpi", cls.pdf_dpi), "device": data.get("runtime", {}).get("device", cls.device), "cache_dir": data.get("runtime", {}).get("cache_dir", ""), "hf_token": data.get("runtime", {}).get("hf_token"), } flat.update(_normalize_overrides(overrides)) allowed = {field.name for field in fields(cls)} return cls(**{key: value for key, value in flat.items() if key in allowed}) @dataclass class MolParserResult: source: str input_index: int page_index: Optional[int] bbox: Optional[tuple[float, float, float, float]] confidence: Optional[float] raw_caption: str caption: str smi: str esmi: str cxsmiles: str markush: bool sru: bool groups: Any def to_dict(self) -> dict[str, Any]: return asdict(self) class MolParser: def __init__(self, config: MolParserConfig | str | Path | None = None, **overrides): if config is None: self.config = MolParserConfig.from_default(**overrides) elif isinstance(config, (str, Path)): self.config = MolParserConfig.from_yaml(config, **overrides) elif overrides: self.config = replace(config, **_normalize_overrides(overrides)) else: self.config = config self._molparser = None self._image_detector = None self._pdf_detector = None def parse( self, inputs, *, rec_only: bool = False, pages=None, expand_px: int | None = None, padding_px: int | None = None, ) -> list[MolParserResult]: with self._temporary_margins(expand_px=expand_px, padding_px=padding_px): image_records: list[tuple[Any, str, int]] = [] pdf_records: list[tuple[Path, str, int]] = [] for input_index, item in enumerate(normalize_inputs(inputs)): if item.kind == "pdf": if item.path is None: raise ValueError("PDF input requires a file path.") pdf_records.append((item.path, item.source, input_index)) else: image_records.append((load_image(item), item.source, input_index)) results: list[MolParserResult] = [] if image_records: results.extend(self._parse_image_batch(image_records, rec_only=rec_only)) if pdf_records: results.extend(self._parse_pdf_batch(pdf_records, pages=pages)) return sorted(results, key=lambda result: result.input_index) def parse_image( self, image, *, rec_only: bool = True, expand_px: int | None = None, padding_px: int | None = None, ) -> list[MolParserResult]: with self._temporary_margins(expand_px=expand_px, padding_px=padding_px): item = normalize_inputs(image)[0] return self._parse_image(load_image(item), item.source, 0, rec_only=rec_only) def parse_pdf( self, pdf, *, pages=None, expand_px: int | None = None, padding_px: int | None = None, ) -> list[MolParserResult]: with self._temporary_margins(expand_px=expand_px, padding_px=padding_px): item = normalize_inputs(pdf)[0] if item.kind != "pdf" or item.path is None: raise ValueError("parse_pdf expects a PDF path or URL.") return self._parse_pdf_batch([(item.path, item.source, 0)], pages=pages) def _token(self): return self.config.hf_token if self.config.hf_token is not None else os.environ.get("HF_TOKEN") def _cache_dir(self): return self.config.cache_dir or None @contextmanager def _temporary_margins(self, *, expand_px: int | None, padding_px: int | None): old_expand = self.config.expand_px old_padding = self.config.padding_px if expand_px is not None: self.config.expand_px = int(expand_px) if padding_px is not None: self.config.padding_px = int(padding_px) try: yield finally: self.config.expand_px = old_expand self.config.padding_px = old_padding def _recognize(self, images: list[Any]) -> list[str]: if self._molparser is None: from .runtime import MolParserRecognizer path = resolve_model( local_path=self.config.molparser_model_path, hf_model_id=self.config.molparser_hf_repo, modelscope_model_id=self.config.molparser_modelscope_repo, cache_dir=self._cache_dir(), token=self._token(), ) self._molparser = MolParserRecognizer( str(path), device=self.config.device, token=self._token(), max_length=self.config.max_length, ) captions: list[str] = [] for batch in _batched(images, self.config.molparser_batch_size): captions.extend(self._molparser.recognize(batch)) return captions def _detector(self, kind: str): from .runtime import MolDetDetector if kind == "pdf": if self._pdf_detector is None: path = self._resolve_detector(self.config.moldet_pdf_model_path, self.config.moldet_pdf_modelname) self._pdf_detector = MolDetDetector( str(path), device=self.config.device, imgsz=self.config.pdf_imgsz, conf=self.config.detector_conf ) return self._pdf_detector if self._image_detector is None: path = self._resolve_detector(self.config.moldet_image_model_path, self.config.moldet_image_modelname) self._image_detector = MolDetDetector( str(path), device=self.config.device, imgsz=self.config.image_imgsz, conf=self.config.detector_conf ) return self._image_detector def _resolve_detector(self, local_path: str, filename: str) -> Path: return resolve_model( local_path=local_path, hf_model_id=self.config.moldet_hf_repo, modelscope_model_id=self.config.moldet_modelscope_repo, filename=filename, cache_dir=self._cache_dir(), token=self._token(), ) def _parse_image(self, image, source: str, input_index: int, *, rec_only: bool) -> list[MolParserResult]: return self._parse_image_batch([(image, source, input_index)], rec_only=rec_only) def _parse_image_batch(self, image_records: list[tuple[Any, str, int]], *, rec_only: bool) -> list[MolParserResult]: images = [record[0] for record in image_records] sources = [record[1] for record in image_records] input_indexes = [record[2] for record in image_records] if rec_only: padded_images = [self._pad_image(image) for image in images] return self._results_many(padded_images, sources, input_indexes, [None] * len(images), [None] * len(images), [None] * len(images)) detections_by_image = self._detect("image", images) crops: list[Any] = [] crop_sources: list[str] = [] crop_input_indexes: list[int] = [] bboxes: list[tuple[float, float, float, float] | None] = [] confidences: list[float | None] = [] for image, source, input_index, detections in zip(images, sources, input_indexes, detections_by_image): for detection in detections: crops.append(self._crop(image, detection)) crop_sources.append(source) crop_input_indexes.append(input_index) bboxes.append(detection.bbox) confidences.append(detection.confidence) return self._results_many( crops, crop_sources, crop_input_indexes, [None] * len(crops), bboxes, confidences, ) def _parse_pdf(self, path: Path, source: str, input_index: int, *, pages=None) -> list[MolParserResult]: return self._parse_pdf_batch([(path, source, input_index)], pages=pages) def _parse_pdf_batch(self, pdf_records: list[tuple[Path, str, int]], *, pages=None) -> list[MolParserResult]: page_records = [] for path, source, input_index in pdf_records: for page in render_pdf(path, dpi=self.config.pdf_dpi, pages=pages): page_records.append((page.image, source, input_index, page.page_index)) if not page_records: return [] page_images = [record[0] for record in page_records] detections_by_page = self._detect("pdf", page_images) crops: list[Any] = [] crop_sources: list[str] = [] crop_input_indexes: list[int] = [] page_indexes: list[int | None] = [] bboxes: list[tuple[float, float, float, float] | None] = [] confidences: list[float | None] = [] for (image, source, input_index, page_index), detections in zip(page_records, detections_by_page): for detection in detections: crops.append(self._crop(image, detection)) crop_sources.append(source) crop_input_indexes.append(input_index) page_indexes.append(page_index) bboxes.append(detection.bbox) confidences.append(detection.confidence) return self._results_many(crops, crop_sources, crop_input_indexes, page_indexes, bboxes, confidences) def _crop(self, image, detection): from .runtime import crop_detection return crop_detection(image, detection, expand_px=self.config.expand_px, pad_px=self.config.padding_px) def _pad_image(self, image): from .runtime import pad_image return pad_image(image, pad_px=self.config.padding_px) def _detect(self, kind: str, images: list[Any]): detector = self._detector(kind) detections: list[Any] = [] for batch in _batched(images, self.config.moldet_batch_size): detections.extend(detector.detect(batch)) return detections def _results( self, images: list[Any], source: str, input_index: int, page_indexes: list[int | None], bboxes: list[tuple[float, float, float, float] | None], confidences: list[float | None], ) -> list[MolParserResult]: return self._results_many( images, [source] * len(images), [input_index] * len(images), page_indexes, bboxes, confidences, ) def _results_many( self, images: list[Any], sources: list[str], input_indexes: list[int], page_indexes: list[int | None], bboxes: list[tuple[float, float, float, float] | None], confidences: list[float | None], ) -> list[MolParserResult]: if not images: return [] raws = self._recognize(images) results: list[MolParserResult] = [] for raw, source, input_index, page_index, bbox, confidence in zip(raws, sources, input_indexes, page_indexes, bboxes, confidences): post = postprocess_caption(raw) results.append( MolParserResult( source=source, input_index=input_index, page_index=page_index, bbox=bbox, confidence=confidence, raw_caption=str(raw), caption=str(post.get("caption", raw)), smi=str(post.get("smi", "")), esmi=str(post.get("esmi", "")), cxsmiles=str(post.get("cxsmiles", "")), markush=bool(post.get("markush", False)), sru=bool(post.get("sru", False)), groups=post.get("groups", ""), ) ) return results def _normalize_overrides(overrides: dict[str, Any]) -> dict[str, Any]: aliases = { "moldet_hf_model_id": "moldet_hf_repo", "moldet_modelscope_model_id": "moldet_modelscope_repo", "recognizer_hf_model_id": "molparser_hf_repo", "recognizer_modelscope_model_id": "molparser_modelscope_repo", "recognizer_model_path": "molparser_model_path", "det_batch_size": "moldet_batch_size", "ocsr_batch_size": "molparser_batch_size", } return {aliases.get(key, key): value for key, value in overrides.items()} def _batched(items: list[Any], batch_size: int): size = max(1, int(batch_size or 1)) for start in range(0, len(items), size): yield items[start : start + size] __all__ = ["MolParser", "MolParserConfig", "MolParserResult"]