368 lines
12 KiB
Python
368 lines
12 KiB
Python
# /pkg/gnu/unitprint/backend/scan/api.py
|
||
import os
|
||
import sys
|
||
import threading
|
||
import io
|
||
import warnings
|
||
from contextlib import asynccontextmanager
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from PIL import Image, ImageOps
|
||
|
||
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
|
||
from fastapi.responses import Response
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
# --- PATHS ---
|
||
SCAN_DIR = Path(__file__).resolve().parent
|
||
DOCSCANNER_DIR = SCAN_DIR / "DocScanner"
|
||
DOCREST_DIR = SCAN_DIR / "DocRes"
|
||
|
||
# Добавляем папки проектов в sys.path, чтобы Python мог найти их внутренние модули.
|
||
# Импорты ниже сработают корректно, так как у проектов нет пересекающихся имен файлов
|
||
# (DocScanner использует `model.py` и `seg.py`, DocRes использует `models/` и `utils.py`).
|
||
for p in (DOCREST_DIR, DOCSCANNER_DIR):
|
||
p_str = str(p)
|
||
if p_str not in sys.path:
|
||
sys.path.insert(0, p_str)
|
||
|
||
# --- DOCSCANNER IMPORTS ---
|
||
from model import DocScanner as DocScannerNet
|
||
from seg import U2NETP
|
||
|
||
# --- DOCRES IMPORTS ---
|
||
from utils import convert_state_dict
|
||
from models import restormer_arch
|
||
from data.preprocess.crop_merge_image import stride_integral
|
||
|
||
|
||
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
USE_HALF = (
|
||
DEVICE.type == "cuda"
|
||
and os.getenv("DOCRES_USE_HALF", "1").strip().lower() not in ("0", "false", "no")
|
||
)
|
||
|
||
API_HOST = os.getenv("API_HOST", "127.0.0.1")
|
||
API_PORT = int(os.getenv("API_PORT", "15200"))
|
||
MAX_SIZE = max(1, int(os.getenv("DOCRES_MAX_SIZE", "1800")))
|
||
|
||
DOCSCANNER_SEG_PATH = DOCSCANNER_DIR / "model_pretrained" / "seg.pth"
|
||
DOCSCANNER_BM_PATH = DOCSCANNER_DIR / "model_pretrained" / "DocScanner-L.pth"
|
||
|
||
_docres_model_env = os.getenv("DOCRES_MODEL_PATH")
|
||
if _docres_model_env:
|
||
DOCRES_MODEL_PATH = Path(_docres_model_env.strip()).expanduser()
|
||
if not DOCRES_MODEL_PATH.is_absolute():
|
||
DOCRES_MODEL_PATH = (SCAN_DIR / DOCRES_MODEL_PATH).resolve()
|
||
else:
|
||
DOCRES_MODEL_PATH = DOCREST_DIR / "data/weights" / "docres.pkl"
|
||
|
||
|
||
_lock = threading.Lock()
|
||
|
||
# =====================================================================
|
||
# --- DOCSCANNER LOGIC ---
|
||
# =====================================================================
|
||
class DocScannerWrapper(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.msk = U2NETP(3, 1)
|
||
self.bm = DocScannerNet()
|
||
|
||
def forward(self, x):
|
||
msk, *_ = self.msk(x)
|
||
bm = self.bm((msk > 0.5).float() * x, iters=12, test_mode=True)
|
||
return (2 * (bm / 286.8) - 1) * 0.99
|
||
|
||
|
||
def load_docscanner_model(model, path, strip_prefix=False):
|
||
state_dict = model.state_dict()
|
||
try:
|
||
pretrained = torch.load(path, map_location=DEVICE, weights_only=False)
|
||
except TypeError:
|
||
pretrained = torch.load(path, map_location=DEVICE)
|
||
|
||
if strip_prefix:
|
||
pretrained = {k[6:]: v for k, v in pretrained.items() if k[6:] in state_dict}
|
||
else:
|
||
pretrained = {k: v for k, v in pretrained.items() if k in state_dict}
|
||
|
||
state_dict.update(pretrained)
|
||
model.load_state_dict(state_dict)
|
||
return model
|
||
|
||
|
||
def process_unwrap(net, payload: bytes) -> bytes:
|
||
try:
|
||
img = Image.open(io.BytesIO(payload))
|
||
try:
|
||
img = ImageOps.exif_transpose(img)
|
||
except Exception:
|
||
pass
|
||
img = img.convert("RGB")
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=400, detail="Uploaded file is not a valid image") from exc
|
||
|
||
image = np.asarray(img, dtype=np.float32) / 255.0
|
||
height, width = image.shape[:2]
|
||
|
||
tensor = (
|
||
torch.from_numpy(cv2.resize(image, (288, 288)).transpose(2, 0, 1))
|
||
.float()
|
||
.unsqueeze(0)
|
||
.to(DEVICE)
|
||
)
|
||
|
||
with _lock, torch.no_grad():
|
||
bm = net(tensor).cpu()
|
||
|
||
flow_x = cv2.blur(cv2.resize(bm[0, 0].numpy(), (width, height)), (3, 3))
|
||
flow_y = cv2.blur(cv2.resize(bm[0, 1].numpy(), (width, height)), (3, 3))
|
||
flow = torch.from_numpy(np.stack([flow_x, flow_y], axis=2)).unsqueeze(0).float()
|
||
|
||
out = F.grid_sample(
|
||
torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float(),
|
||
flow,
|
||
align_corners=True,
|
||
)
|
||
|
||
result = (out[0] * 255).permute(1, 2, 0).numpy()
|
||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||
result_bgr = np.ascontiguousarray(result[:, :, ::-1])
|
||
|
||
ok, buf = cv2.imencode(".png", result_bgr)
|
||
if not ok:
|
||
raise HTTPException(status_code=500, detail="Failed to encode result image")
|
||
return buf.tobytes()
|
||
|
||
|
||
# =====================================================================
|
||
# --- DOCRES LOGIC ---
|
||
# =====================================================================
|
||
def load_docres_model(model, path):
|
||
try:
|
||
ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
|
||
except TypeError:
|
||
ckpt = torch.load(path, map_location=DEVICE)
|
||
|
||
if isinstance(ckpt, dict):
|
||
state = None
|
||
for key in ("model_state", "state_dict", "model", "params", "params_ema"):
|
||
if key in ckpt and isinstance(ckpt[key], dict):
|
||
state = ckpt[key]
|
||
break
|
||
if state is None:
|
||
state = ckpt
|
||
else:
|
||
raise RuntimeError(f"Unsupported checkpoint format: {path}")
|
||
|
||
state = convert_state_dict(state)
|
||
model.load_state_dict(state)
|
||
model.eval()
|
||
model.to(DEVICE)
|
||
|
||
if USE_HALF:
|
||
model.half()
|
||
else:
|
||
model.float()
|
||
return model
|
||
|
||
|
||
def resize_if_needed(img: np.ndarray, max_size: int = MAX_SIZE) -> np.ndarray:
|
||
h, w = img.shape[:2]
|
||
longest = max(h, w)
|
||
if longest <= 0:
|
||
raise HTTPException(status_code=400, detail="Image has zero size")
|
||
scale = min(max_size / longest, 1.0)
|
||
if scale < 1.0:
|
||
new_w = max(1, int(w * scale))
|
||
new_h = max(1, int(h * scale))
|
||
return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||
return img
|
||
|
||
|
||
def deblur_prompt(img: np.ndarray) -> np.ndarray:
|
||
x = cv2.Sobel(img, cv2.CV_16S, 1, 0)
|
||
y = cv2.Sobel(img, cv2.CV_16S, 0, 1)
|
||
absX = cv2.convertScaleAbs(x)
|
||
absY = cv2.convertScaleAbs(y)
|
||
high_frequency = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
|
||
high_frequency = cv2.cvtColor(high_frequency, cv2.COLOR_BGR2GRAY)
|
||
high_frequency = cv2.cvtColor(high_frequency, cv2.COLOR_GRAY2BGR)
|
||
return high_frequency
|
||
|
||
|
||
def appearance_prompt(img: np.ndarray) -> np.ndarray:
|
||
h, w = img.shape[:2]
|
||
img_1024 = cv2.resize(img, (1024, 1024))
|
||
planes = cv2.split(img_1024)
|
||
norm_planes = []
|
||
for plane in planes:
|
||
dilated_img = cv2.dilate(plane, np.ones((7, 7), np.uint8))
|
||
bg_img = cv2.medianBlur(dilated_img, 21)
|
||
diff_img = 255 - cv2.absdiff(plane, bg_img)
|
||
norm_img = cv2.normalize(diff_img, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
|
||
norm_planes.append(norm_img)
|
||
result_norm = cv2.merge(norm_planes)
|
||
result_norm = cv2.resize(result_norm, (w, h))
|
||
return result_norm
|
||
|
||
|
||
def _run_restormer(model, in_im: np.ndarray, padding_h: int, padding_w: int) -> np.ndarray:
|
||
arr = in_im.astype(np.float32) / 255.0
|
||
with _lock, torch.no_grad():
|
||
tensor = torch.from_numpy(arr.transpose(2, 0, 1)).unsqueeze(0).to(DEVICE)
|
||
tensor = tensor.half() if USE_HALF else tensor.float()
|
||
pred = model(tensor)
|
||
pred = torch.clamp(pred, 0, 1)
|
||
pred = pred[0].permute(1, 2, 0).cpu().numpy()
|
||
pred = (pred * 255).astype(np.uint8)
|
||
out = pred[padding_h:, padding_w:]
|
||
if out.size == 0:
|
||
raise HTTPException(status_code=500, detail="Empty result after padding crop")
|
||
return out
|
||
|
||
|
||
def deblur_image(model, img: np.ndarray) -> np.ndarray:
|
||
img = resize_if_needed(img)
|
||
in_im, padding_h, padding_w = stride_integral(img, 8)
|
||
prompt = deblur_prompt(in_im)
|
||
in_im = np.concatenate((in_im, prompt), -1)
|
||
return _run_restormer(model, in_im, padding_h, padding_w)
|
||
|
||
|
||
def appearance_image(model, img: np.ndarray) -> np.ndarray:
|
||
img = resize_if_needed(img)
|
||
prompt = appearance_prompt(img)
|
||
in_im = np.concatenate((img, prompt), -1)
|
||
in_im, padding_h, padding_w = stride_integral(in_im, 8)
|
||
return _run_restormer(model, in_im, padding_h, padding_w)
|
||
|
||
|
||
def process_docres(model, payload: bytes, task: str) -> bytes:
|
||
img = cv2.imdecode(np.frombuffer(payload, np.uint8), cv2.IMREAD_COLOR)
|
||
if img is None or img.size == 0:
|
||
raise HTTPException(status_code=400, detail="Uploaded file is not a valid image")
|
||
|
||
try:
|
||
if task == "deblur":
|
||
result = deblur_image(model, img)
|
||
elif task == "illumination_correct":
|
||
result = appearance_image(model, img)
|
||
else:
|
||
raise HTTPException(status_code=500, detail=f"Unknown task: {task}")
|
||
except HTTPException:
|
||
raise
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=500, detail=f"{task} failed: {exc}") from exc
|
||
|
||
ok, buf = cv2.imencode(".png", np.ascontiguousarray(result))
|
||
if not ok:
|
||
raise HTTPException(status_code=500, detail="Failed to encode result image")
|
||
return buf.tobytes()
|
||
|
||
|
||
# =====================================================================
|
||
# --- FASTAPI APP ---
|
||
# =====================================================================
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
if not DOCSCANNER_SEG_PATH.exists():
|
||
raise RuntimeError(f"DocScanner seg.pth not found: {DOCSCANNER_SEG_PATH}")
|
||
if not DOCSCANNER_BM_PATH.exists():
|
||
raise RuntimeError(f"DocScanner DocScanner-L.pth not found: {DOCSCANNER_BM_PATH}")
|
||
if not DOCRES_MODEL_PATH.exists():
|
||
raise RuntimeError(f"DocRes checkpoint not found: {DOCRES_MODEL_PATH}")
|
||
|
||
# Init DocScanner
|
||
docscanner_net = DocScannerWrapper().to(DEVICE).eval()
|
||
load_docscanner_model(docscanner_net.msk, DOCSCANNER_SEG_PATH, strip_prefix=True)
|
||
load_docscanner_model(docscanner_net.bm, DOCSCANNER_BM_PATH, strip_prefix=False)
|
||
|
||
# Init DocRes
|
||
docres_model = restormer_arch.Restormer(
|
||
inp_channels=6, out_channels=3, dim=48,
|
||
num_blocks=[2, 3, 3, 4], num_refinement_blocks=4,
|
||
heads=[1, 2, 4, 8], ffn_expansion_factor=2.66,
|
||
bias=False, LayerNorm_type="WithBias", dual_pixel_task=True,
|
||
)
|
||
load_docres_model(docres_model, DOCRES_MODEL_PATH)
|
||
|
||
app.state.docscanner_net = docscanner_net
|
||
app.state.docres_model = docres_model
|
||
yield
|
||
|
||
|
||
app = FastAPI(title="Scan Backend API", lifespan=lifespan)
|
||
|
||
|
||
@app.get("/health")
|
||
def health():
|
||
return {
|
||
"status": "ok",
|
||
"device": str(DEVICE),
|
||
"docres_half_precision": USE_HALF,
|
||
"docres_model_path": str(DOCRES_MODEL_PATH),
|
||
"docscanner_paths": {
|
||
"seg": str(DOCSCANNER_SEG_PATH),
|
||
"bm": str(DOCSCANNER_BM_PATH)
|
||
}
|
||
}
|
||
|
||
|
||
@app.post("/unwrap")
|
||
def unwrap(request: Request, file: UploadFile = File(...)):
|
||
try:
|
||
file.file.seek(0)
|
||
payload = file.file.read()
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=400, detail="Cannot read uploaded file") from exc
|
||
if not payload:
|
||
raise HTTPException(status_code=400, detail="Uploaded file is empty")
|
||
|
||
result_png = process_unwrap(request.app.state.docscanner_net, payload)
|
||
return Response(
|
||
content=result_png,
|
||
media_type="image/png",
|
||
headers={"Content-Disposition": 'inline; filename="unwrap.png"', "Cache-Control": "no-store"},
|
||
)
|
||
|
||
|
||
def _docres_endpoint(request: Request, file: UploadFile, task: str):
|
||
try:
|
||
file.file.seek(0)
|
||
payload = file.file.read()
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=400, detail="Cannot read uploaded file") from exc
|
||
if not payload:
|
||
raise HTTPException(status_code=400, detail="Uploaded file is empty")
|
||
|
||
result_png = process_docres(request.app.state.docres_model, payload, task)
|
||
return Response(
|
||
content=result_png,
|
||
media_type="image/png",
|
||
headers={"Content-Disposition": f'inline; filename="{task}.png"', "Cache-Control": "no-store"},
|
||
)
|
||
|
||
|
||
@app.post("/deblur")
|
||
def deblur(request: Request, file: UploadFile = File(...)):
|
||
return _docres_endpoint(request, file, "deblur")
|
||
|
||
|
||
@app.post("/illumination_correct")
|
||
def illumination_correct(request: Request, file: UploadFile = File(...)):
|
||
return _docres_endpoint(request, file, "illumination_correct")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
uvicorn.run(app, host=API_HOST, port=API_PORT, workers=1)
|