app.py

# -*- coding: utf-8 -*-
import base64
import binascii
import io
import json
import os
import tempfile
import time
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr
import numpy as np
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from PIL import Image

from hf_loader import model_descriptor, predict_image_file

app = FastAPI(
    title="MolScribe OCR",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

AUTO_TRIM_WHITE = str(os.getenv("AUTO_TRIM_WHITE", "1")).strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}
WHITE_THRESHOLD = int(os.getenv("WHITE_THRESHOLD", "245") or 245)
WHITE_PADDING = max(0, int(os.getenv("WHITE_PADDING", "16") or 16))
MAX_IMAGE_EDGE = max(0, int(os.getenv("MAX_IMAGE_EDGE", "1280") or 1280))
MIN_IMAGE_EDGE = max(0, int(os.getenv("MIN_IMAGE_EDGE", "0") or 0))


class MolScribeRequest(BaseModel):
    image_base64: str = Field(
        ...,
        description="Base64 编码的图片内容，支持纯 base64 字符串或 data URL。",
    )
    return_atoms_bonds: Optional[bool] = True
    return_confidence: Optional[bool] = True
    timeout_seconds: Optional[float] = Field(
        default=None,
        ge=0,
        description="可选请求超时秒数；0 或 null 表示使用服务端默认值。",
    )


class BatchRequest(BaseModel):
    inputs: List[MolScribeRequest]


def _to_jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): _to_jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_jsonable(v) for v in value]
    if hasattr(value, "item") and callable(value.item):
        try:
            return value.item()
        except Exception:
            return str(value)
    return value


def _decode_image_bytes(image_base64: str) -> bytes:
    payload = (image_base64 or "").strip()
    if not payload:
        raise ValueError("image_base64 为空")

    if payload.startswith("data:"):
        parts = payload.split(",", 1)
        if len(parts) != 2:
            raise ValueError("data URL 格式不正确")
        payload = parts[1]

    try:
        raw = base64.b64decode(payload, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise ValueError("image_base64 不是合法的 Base64 图片数据") from exc

    if not raw:
        raise ValueError("解码后的图片为空")
    return raw


def _white_content_bbox(image: Image.Image) -> Optional[Tuple[int, int, int, int]]:
    rgb = image.convert("RGB")
    arr = np.asarray(rgb)
    if arr.ndim != 3 or arr.shape[2] < 3:
        return None
    mask = np.any(arr < WHITE_THRESHOLD, axis=2)
    if not bool(mask.any()):
        return None
    coords = np.argwhere(mask)
    y0, x0 = coords.min(axis=0)
    y1, x1 = coords.max(axis=0) + 1
    return int(x0), int(y0), int(x1), int(y1)


def _trim_white_border(image: Image.Image) -> Tuple[Image.Image, Dict[str, Any]]:
    bbox = _white_content_bbox(image)
    width, height = image.size
    if bbox is None:
        return image, {"trimmed": False, "trim_bbox": [0, 0, width, height]}

    x0, y0, x1, y1 = bbox
    left = max(0, x0 - WHITE_PADDING)
    top = max(0, y0 - WHITE_PADDING)
    right = min(width, x1 + WHITE_PADDING)
    bottom = min(height, y1 + WHITE_PADDING)

    trimmed = image.crop((left, top, right, bottom))
    changed = (left, top, right, bottom) != (0, 0, width, height)
    return trimmed, {
        "trimmed": bool(changed),
        "trim_bbox": [int(left), int(top), int(right), int(bottom)],
    }


def _auto_scale_image(image: Image.Image) -> Tuple[Image.Image, Dict[str, Any]]:
    width, height = image.size
    max_edge = max(width, height)
    if max_edge <= 0:
        return image, {"scaled": False, "scale_factor": 1.0}

    scale_factor = 1.0
    if MAX_IMAGE_EDGE > 0 and max_edge > MAX_IMAGE_EDGE:
        scale_factor = min(scale_factor, MAX_IMAGE_EDGE / float(max_edge))
    if MIN_IMAGE_EDGE > 0 and max_edge < MIN_IMAGE_EDGE:
        scale_factor = max(scale_factor, MIN_IMAGE_EDGE / float(max_edge))

    if abs(scale_factor - 1.0) < 1e-6:
        return image, {"scaled": False, "scale_factor": 1.0}

    target_size = (
        max(1, int(round(width * scale_factor))),
        max(1, int(round(height * scale_factor))),
    )
    resized = image.resize(target_size, Image.Resampling.LANCZOS)
    return resized, {
        "scaled": True,
        "scale_factor": round(float(scale_factor), 4),
    }


def _write_preprocessed_png(image: Image.Image) -> Tuple[str, Dict[str, Any]]:
    normalized = image.convert("RGB")
    original_width, original_height = normalized.size

    processed = normalized
    trim_info = {
        "trimmed": False,
        "trim_bbox": [0, 0, int(original_width), int(original_height)],
    }
    if AUTO_TRIM_WHITE:
        processed, trim_info = _trim_white_border(processed)

    processed, scale_info = _auto_scale_image(processed)

    image_info = {
        "mode": "RGB",
        "original_width": int(original_width),
        "original_height": int(original_height),
        "width": int(processed.width),
        "height": int(processed.height),
        "trimmed": trim_info["trimmed"],
        "trim_bbox": trim_info["trim_bbox"],
        "scaled": scale_info["scaled"],
        "scale_factor": scale_info["scale_factor"],
    }

    fd, temp_path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    processed.save(temp_path, format="PNG")
    return temp_path, image_info


def _write_png_from_base64(image_base64: str) -> Tuple[str, Dict[str, Any]]:
    raw = _decode_image_bytes(image_base64)
    try:
        with Image.open(io.BytesIO(raw)) as image:
            return _write_preprocessed_png(image)
    except Exception as exc:
        raise ValueError("无法将输入内容解析为图片") from exc


def _write_png_from_upload(image_path: str) -> Tuple[str, Dict[str, Any]]:
    try:
        with Image.open(image_path) as image:
            return _write_preprocessed_png(image)
    except Exception as exc:
        raise ValueError("无法读取上传图片") from exc


def _predict_request(req: MolScribeRequest) -> Dict[str, Any]:
    started = time.perf_counter()
    temp_path = None

    try:
        temp_path, image_info = _write_png_from_base64(req.image_base64)
        prediction = predict_image_file(
            temp_path,
            return_atoms_bonds=req.return_atoms_bonds
            if req.return_atoms_bonds is not None
            else True,
            return_confidence=req.return_confidence
            if req.return_confidence is not None
            else True,
            timeout_seconds=req.timeout_seconds,
        )
        safe_prediction = _to_jsonable(prediction)
        return {
            "success": True,
            "smiles": safe_prediction.get("smiles", ""),
            "prediction": safe_prediction,
            "image": image_info,
            "model": model_descriptor(),
            "elapsed_ms": round((time.perf_counter() - started) * 1000, 2),
        }
    except Exception as exc:
        return {
            "success": False,
            "error": str(exc),
        }
    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)


def _predict_uploaded_file(
    image_path: str,
    return_atoms_bonds: bool,
    return_confidence: bool,
    timeout_seconds: float | None = None,
) -> Dict[str, Any]:
    if not image_path:
        return {"success": False, "error": "未上传图片"}

    started = time.perf_counter()
    temp_path = None
    try:
        temp_path, image_info = _write_png_from_upload(image_path)
        prediction = predict_image_file(
            temp_path,
            return_atoms_bonds=return_atoms_bonds,
            return_confidence=return_confidence,
            timeout_seconds=timeout_seconds,
        )
        safe_prediction = _to_jsonable(prediction)
        return {
            "success": True,
            "smiles": safe_prediction.get("smiles", ""),
            "prediction": safe_prediction,
            "image": image_info,
            "model": model_descriptor(),
            "elapsed_ms": round((time.perf_counter() - started) * 1000, 2),
        }
    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)


@app.get("/healthz")
def healthz():
    return {
        "ok": True,
        "model": model_descriptor(),
    }


@app.post("/api/molscribe")
def api_molscribe(req: MolScribeRequest):
    return _predict_request(req)


@app.post("/api/molscribe/batch")
def api_molscribe_batch(req: BatchRequest):
    return [_predict_request(item) for item in req.inputs]


def gradio_fn(image_path: str, return_atoms_bonds: bool, return_confidence: bool):
    try:
        result = _predict_uploaded_file(
            image_path,
            return_atoms_bonds,
            return_confidence,
            timeout_seconds=None,
        )
        if not result.get("success"):
            return "", "", json.dumps(result, ensure_ascii=False, indent=2)
        prediction = result["prediction"]
        return (
            result.get("smiles", ""),
            prediction.get("molfile", ""),
            json.dumps(result, ensure_ascii=False, indent=2),
        )
    except Exception as exc:
        payload = {"success": False, "error": str(exc)}
        return "", "", json.dumps(payload, ensure_ascii=False, indent=2)


demo = gr.Interface(
    fn=gradio_fn,
    inputs=[
        gr.Image(type="filepath", label="上传化学结构图片"),
        gr.Checkbox(label="返回 atoms / bonds", value=True),
        gr.Checkbox(label="返回 confidence", value=True),
    ],
    outputs=[
        gr.Textbox(label="识别出的 SMILES", interactive=True),
        gr.Textbox(label="Molfile", lines=14, interactive=True),
        gr.Textbox(label="完整 JSON 结果", lines=20),
    ],
    title="MolScribe OCR",
    description="上传化学结构图片，返回 MolScribe 模型识别出的 SMILES、molfile、confidence、atoms 和 bonds。",
)

app = gr.mount_gradio_app(app, demo, path="/")