update

2026-05-14 13:30:06 +08:00
parent 974946cee4
commit e43171521d
91 changed files with 10485 additions and 3254 deletions
@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+
+DEFAULT_LAMA_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+
+
+def _default_cache_path() -> Path:
+    # 放在 python_server/outputs/cache 下，避免污染用户家目录
+    here = Path(__file__).resolve()
+    cache_dir = here.parents[2] / "outputs" / "cache"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir / "big-lama.pt"
+
+
+def _download(url: str, dst: Path) -> None:
+    import requests
+
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    resp = requests.get(url, stream=True, timeout=120)
+    resp.raise_for_status()
+    tmp = dst.with_suffix(dst.suffix + ".tmp")
+    with tmp.open("wb") as f:
+        for chunk in resp.iter_content(chunk_size=1024 * 1024):
+            if not chunk:
+                continue
+            f.write(chunk)
+    tmp.replace(dst)
+
+
+def _pad_to_mod8(img: np.ndarray) -> tuple[np.ndarray, tuple[int, int]]:
+    h, w = img.shape[:2]
+    pad_h = (8 - (h % 8)) % 8
+    pad_w = (8 - (w % 8)) % 8
+    if pad_h == 0 and pad_w == 0:
+        return img, (0, 0)
+    out = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
+    return out, (pad_h, pad_w)
+
+
+def _unpad(img: np.ndarray, pad_hw: tuple[int, int]) -> np.ndarray:
+    pad_h, pad_w = pad_hw
+    if pad_h == 0 and pad_w == 0:
+        return img
+    h, w = img.shape[:2]
+    return img[: h - pad_h, : w - pad_w]
+
+
+class LaMaTorchscript:
+    """
+    轻量 LaMa（big-lama.pt torchscript）推理封装。
+    - 输入：RGB 图 + L mask（白=需要修补）
+    - 输出：RGB 图（同尺寸）
+    """
+
+    def __init__(self, device: str = "cpu", model_path: Optional[str] = None, model_url: Optional[str] = None):
+        import torch
+
+        self.device = device
+        self.model_url = model_url or DEFAULT_LAMA_URL
+        p = Path(model_path) if model_path else _default_cache_path()
+        if not p.exists():
+            _download(self.model_url, p)
+        self.model_path = p
+        self.model = torch.jit.load(str(self.model_path), map_location=device).eval()
+
+    @staticmethod
+    def _norm_img_u8(rgb: np.ndarray) -> np.ndarray:
+        # [H,W,3] uint8 -> float32 [0,1]
+        return (rgb.astype(np.float32) / 255.0).clip(0.0, 1.0)
+
+    def __call__(self, image: Image.Image, mask: Image.Image) -> Image.Image:
+        import torch
+
+        img = np.asarray(image.convert("RGB"), dtype=np.uint8)
+        m = np.asarray(mask.convert("L"), dtype=np.uint8)
+        m = (m >= 128).astype(np.float32)  # 1=hole
+
+        img_f = self._norm_img_u8(img)
+        # pad to multiple of 8
+        img_pad, pad_hw = _pad_to_mod8(img_f)
+        m_pad, _ = _pad_to_mod8(np.repeat(m[:, :, None], 3, axis=2))
+        m1 = m_pad[:, :, 0]  # [H,W]
+
+        # torch: [1,C,H,W]
+        it = torch.from_numpy(img_pad).permute(2, 0, 1).unsqueeze(0).to(self.device)
+        mt = torch.from_numpy(m1).unsqueeze(0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            out = self.model(it, mt)
+
+        out_np = out[0].permute(1, 2, 0).detach().cpu().numpy()
+        out_np = np.clip(out_np * 255.0, 0.0, 255.0).astype(np.uint8)
+        out_np = _unpad(out_np, pad_hw)
+        return Image.fromarray(out_np, mode="RGB")
+