From 1d71f21671672d23c63e02db8f7b43d9318d2d6c Mon Sep 17 00:00:00 2001 From: TraceLumen Date: Mon, 14 Sep 2026 06:09:13 +0300 Subject: [PATCH] init --- .gitignore | 3 + api.py | 367 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 370 insertions(+) create mode 100644 .gitignore create mode 100644 api.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2722a76 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +DocRes +DocScanner +venv diff --git a/api.py b/api.py new file mode 100644 index 0000000..bd89034 --- /dev/null +++ b/api.py @@ -0,0 +1,367 @@ +# /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 / "checkpoints" / "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)