106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
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")
|
||
|