from __future__ import annotations import contextlib from dataclasses import dataclass from pathlib import Path from typing import Optional, Sequence from PIL import Image @dataclass class ModelResolveError(RuntimeError): message: str def __str__(self) -> str: return self.message @dataclass class Detection: bbox: tuple[float, float, float, float] confidence: float def resolve_model( *, local_path: str | Path | None = None, hf_model_id: str = "", modelscope_model_id: str = "", filename: str | None = None, cache_dir: str | Path | None = None, token: str | bool | None = None, ) -> Path: existing = _as_existing_path(local_path) if existing is not None: return _find_file(existing, filename) if filename and existing.is_dir() else existing errors: list[str] = [] if hf_model_id: try: if filename: from huggingface_hub import hf_hub_download return Path( hf_hub_download( repo_id=hf_model_id, filename=filename, cache_dir=str(cache_dir) if cache_dir else None, token=token, ) ) from huggingface_hub import snapshot_download return Path(snapshot_download(repo_id=hf_model_id, cache_dir=str(cache_dir) if cache_dir else None, token=token)) except Exception as exc: errors.append(f"Hugging Face failed for {hf_model_id}: {exc!r}") if modelscope_model_id: try: try: from modelscope import snapshot_download except ImportError: from modelscope.hub.snapshot_download import snapshot_download root = Path(snapshot_download(modelscope_model_id, cache_dir=str(cache_dir) if cache_dir else None)) return _find_file(root, filename) if filename else root except Exception as exc: errors.append(f"ModelScope failed for {modelscope_model_id}: {exc!r}") details = " ".join(errors) if errors else "No model id was configured." raise ModelResolveError(f"Unable to resolve model. {details} Set a local path or configure HF/ModelScope.") class MolDetDetector: def __init__(self, model_path: str, *, device: str = "auto", imgsz: int = 640, conf: float = 0.5): try: from ultralytics import YOLO except ImportError as exc: raise ImportError("MolDet requires ultralytics. Install molparser[moldet] or molparser[inference].") from exc self.model = YOLO(str(model_path)) self.device = None if device == "auto" else device self.imgsz = int(imgsz) self.conf = float(conf) def detect(self, images: Sequence[Image.Image]) -> list[list[Detection]]: results = self.model.predict(list(images), imgsz=self.imgsz, conf=self.conf, device=self.device, verbose=False) batches: list[list[Detection]] = [] for result in results: detections: list[Detection] = [] boxes = getattr(result, "boxes", None) if boxes is not None: for box in boxes: xyxy = tuple(float(v) for v in box.xyxy[0].tolist()) detections.append(Detection(bbox=xyxy, confidence=float(box.conf))) batches.append(detections) return batches class MolParserRecognizer: def __init__( self, model_path: str, *, device: str = "auto", token=None, max_length: int = 256, ): try: import torch from transformers import AutoModelForImageTextToText, AutoProcessor except ImportError as exc: raise ImportError("MolParser recognition requires torch and transformers. Install molparser[inference].") from exc self.torch = torch self.device = torch.device("cuda" if device == "auto" and torch.cuda.is_available() else ("cpu" if device == "auto" else device)) self.dtype = torch.float16 if self.device.type == "cuda" else torch.float32 self.max_length = int(max_length) self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True, token=token) self.model = AutoModelForImageTextToText.from_pretrained( model_path, trust_remote_code=True, token=token, dtype=self.dtype, ).to(self.device) self.model.eval() def recognize(self, images: Sequence[Image.Image]) -> list[str]: inputs = self.processor(images=[image.convert("RGB") for image in images], return_tensors="pt") pixel_values = inputs["pixel_values"].to(device=self.device, dtype=self.dtype) autocast = ( self.torch.autocast(device_type="cuda", dtype=self.dtype) if self.device.type == "cuda" else contextlib.nullcontext() ) with self.torch.inference_mode(), autocast: outputs = self.model.generate( pixel_values=pixel_values, max_length=self.max_length, num_beams=1, do_sample=False, use_cache=True, ) return self.processor.batch_decode(outputs.detach().cpu().tolist(), skip_special_tokens=True) def crop_detection(image: Image.Image, detection: Detection, *, expand_px: int = 0, pad_px: int = 0) -> Image.Image: width, height = image.size x1, y1, x2, y2 = detection.bbox box = ( max(0, int(round(x1)) - expand_px), max(0, int(round(y1)) - expand_px), min(width, int(round(x2)) + expand_px), min(height, int(round(y2)) + expand_px), ) crop = image.crop(box).convert("RGB") return pad_image(crop, pad_px=pad_px) def pad_image(image: Image.Image, *, pad_px: int = 0) -> Image.Image: image = image.convert("RGB") if pad_px <= 0: return image padded = Image.new("RGB", (image.width + pad_px * 2, image.height + pad_px * 2), "white") padded.paste(image, (pad_px, pad_px)) return padded def _as_existing_path(path: str | Path | None) -> Optional[Path]: if not path: return None resolved = Path(path).expanduser() return resolved if resolved.exists() else None def _find_file(root: Path, filename: str | None) -> Path: if not filename: return root candidate = root / filename if candidate.exists(): return candidate matches = list(root.rglob(filename)) if matches: return matches[0] raise ModelResolveError(f"{filename} not found under model directory: {root}") MobileRecognizer = MolParserRecognizer __all__ = [ "Detection", "MobileRecognizer", "ModelResolveError", "MolDetDetector", "MolParserRecognizer", "crop_detection", "pad_image", "resolve_model", ]