flakego commited on
Commit
288bb46
·
verified ·
1 Parent(s): 84d8db2

Lazy-load torch so health starts without model import

Browse files
Files changed (1) hide show
  1. app.py +584 -580
app.py CHANGED
@@ -1,580 +1,584 @@
1
- from __future__ import annotations
2
-
3
- import hashlib
4
- import io
5
- import os
6
- import threading
7
- import time
8
- from pathlib import Path
9
- from urllib.parse import urlparse
10
-
11
- import cv2
12
- import numpy as np
13
- import onnxruntime as ort
14
- import torch
15
- from fastapi import FastAPI, File, Form, HTTPException, UploadFile
16
- from fastapi.responses import Response
17
- from huggingface_hub import hf_hub_download
18
- from PIL import Image
19
- from torch.hub import download_url_to_file, get_dir
20
-
21
-
22
- LAMA_MODEL_REPO = os.getenv("LAMA_MODEL_REPO", "Carve/LaMa-ONNX")
23
- LAMA_MODEL_FILE = os.getenv("LAMA_MODEL_FILE", "lama_fp32.onnx")
24
- MIGAN_MODEL_URL = os.getenv(
25
- "MIGAN_MODEL_URL",
26
- "https://github.com/Sanster/models/releases/download/migan/migan_traced.pt",
27
- )
28
- MIGAN_MODEL_MD5 = os.getenv("MIGAN_MODEL_MD5", "76eb3b1a71c400ee3290524f7a11b89c")
29
- DEFAULT_MODE = os.getenv("DEFAULT_MODE", "quality")
30
- LOAD_MODEL_ON_STARTUP = os.getenv("LOAD_MODEL_ON_STARTUP", "1") == "1"
31
- PROVIDERS = ["CPUExecutionProvider"]
32
- MODES = {
33
- "quality": "selfbuilt_lama_onnx",
34
- "fast": "iopaint_migan_torchscript",
35
- }
36
-
37
- DOUBAO_ALPHA_ASSET = Path(__file__).with_name("doubao_alpha.png")
38
- DOUBAO_WIDTH_FRAC = 0.22
39
- DOUBAO_HEIGHT_FRAC = 0.075
40
- DOUBAO_MARGIN_RIGHT_FRAC = 0.004
41
- DOUBAO_MARGIN_BOTTOM_FRAC = 0.004
42
- DOUBAO_MAX_SATURATION = 55
43
- DOUBAO_LOGO_MIN_LUMA = 150
44
- DOUBAO_TOPHAT_DELTA = 12
45
- DOUBAO_DETECT_MIN_COVERAGE = 0.04
46
- DOUBAO_DETECT_NCC_THRESHOLD = 0.4
47
- DOUBAO_ALPHA_WIDTH_FRAC = 0.1636
48
- DOUBAO_ALPHA_HEIGHT_FRAC = 0.0405
49
- DOUBAO_ALPHA_MARGIN_RIGHT_FRAC = 0.0132
50
- DOUBAO_ALPHA_MARGIN_BOTTOM_FRAC = 0.0166
51
- DOUBAO_ALPHA_ALIGN_SEARCH = (0.88, 1.12, 25)
52
- DOUBAO_RESIDUAL_ALPHA_FLOOR = 0.05
53
- DOUBAO_TEMPLATE_DILATE = 4
54
-
55
- LAMA_SESSION: ort.InferenceSession | None = None
56
- MIGAN_MODEL: torch.jit.ScriptModule | None = None
57
- DOUBAO_ALPHA: np.ndarray | None = None
58
- DOUBAO_SILHOUETTE: np.ndarray | None = None
59
- MODEL_LOCK = threading.RLock()
60
-
61
- app = FastAPI(title="Image Services Inpaint")
62
-
63
-
64
- def _md5sum(path: Path) -> str:
65
- digest = hashlib.md5()
66
- with path.open("rb") as handle:
67
- for chunk in iter(lambda: handle.read(1024 * 1024), b""):
68
- digest.update(chunk)
69
- return digest.hexdigest()
70
-
71
-
72
- def _cache_path_by_url(url: str) -> Path:
73
- parts = urlparse(url)
74
- model_dir = Path(get_dir()) / "checkpoints"
75
- model_dir.mkdir(parents=True, exist_ok=True)
76
- return model_dir / Path(parts.path).name
77
-
78
-
79
- def _download_torchscript(url: str, expected_md5: str) -> Path:
80
- path = _cache_path_by_url(url)
81
- if not path.exists():
82
- download_url_to_file(url, str(path), None, progress=True)
83
- actual_md5 = _md5sum(path)
84
- if actual_md5 != expected_md5:
85
- try:
86
- path.unlink()
87
- finally:
88
- raise RuntimeError(f"Model md5 mismatch: {actual_md5} != {expected_md5}")
89
- return path
90
-
91
-
92
- def _load_lama_session() -> ort.InferenceSession:
93
- global LAMA_SESSION
94
- if LAMA_SESSION is not None:
95
- return LAMA_SESSION
96
- with MODEL_LOCK:
97
- if LAMA_SESSION is not None:
98
- return LAMA_SESSION
99
- model_path = hf_hub_download(repo_id=LAMA_MODEL_REPO, filename=LAMA_MODEL_FILE)
100
- LAMA_SESSION = ort.InferenceSession(model_path, providers=PROVIDERS)
101
- return LAMA_SESSION
102
-
103
-
104
- def _load_migan_model() -> torch.jit.ScriptModule:
105
- global MIGAN_MODEL
106
- if MIGAN_MODEL is not None:
107
- return MIGAN_MODEL
108
- with MODEL_LOCK:
109
- if MIGAN_MODEL is not None:
110
- return MIGAN_MODEL
111
- model_path = _download_torchscript(MIGAN_MODEL_URL, MIGAN_MODEL_MD5)
112
- MIGAN_MODEL = torch.jit.load(str(model_path), map_location="cpu").eval()
113
- return MIGAN_MODEL
114
-
115
-
116
- def _prepare_mask(mask: Image.Image, size: tuple[int, int]) -> np.ndarray:
117
- arr = np.array(mask.convert("L").resize(size, Image.Resampling.NEAREST), copy=True)
118
- return np.where(arr > 127, 255, 0).astype(np.uint8)
119
-
120
-
121
- def _load_doubao_alpha() -> np.ndarray:
122
- global DOUBAO_ALPHA
123
- if DOUBAO_ALPHA is not None:
124
- return DOUBAO_ALPHA
125
- alpha = cv2.imread(str(DOUBAO_ALPHA_ASSET), cv2.IMREAD_GRAYSCALE)
126
- if alpha is None:
127
- raise RuntimeError(f"Missing Doubao alpha asset: {DOUBAO_ALPHA_ASSET}")
128
- DOUBAO_ALPHA = alpha.astype(np.float32) / 255.0
129
- return DOUBAO_ALPHA
130
-
131
-
132
- def _load_doubao_silhouette() -> np.ndarray:
133
- global DOUBAO_SILHOUETTE
134
- if DOUBAO_SILHOUETTE is not None:
135
- return DOUBAO_SILHOUETTE
136
- DOUBAO_SILHOUETTE = (_load_doubao_alpha() > 0.15).astype(np.uint8) * 255
137
- return DOUBAO_SILHOUETTE
138
-
139
-
140
- def _doubao_locate(image_bgr: np.ndarray) -> tuple[int, int, int, int]:
141
- height, width = image_bgr.shape[:2]
142
- mark_w = max(40, int(width * DOUBAO_WIDTH_FRAC))
143
- mark_h = max(16, int(width * DOUBAO_HEIGHT_FRAC))
144
- margin_right = max(4, int(width * DOUBAO_MARGIN_RIGHT_FRAC))
145
- margin_bottom = max(4, int(width * DOUBAO_MARGIN_BOTTOM_FRAC))
146
- x = max(0, width - margin_right - mark_w)
147
- y = max(0, height - margin_bottom - mark_h)
148
- return x, y, min(mark_w, width - x), min(mark_h, height - y)
149
-
150
-
151
- def _doubao_extract_candidate_mask(image_bgr: np.ndarray, loc: tuple[int, int, int, int]) -> np.ndarray:
152
- x, y, width, height = loc
153
- if width < 16 or height < 16:
154
- return np.zeros((height, width), np.uint8)
155
- roi = image_bgr[y : y + height, x : x + width].astype(np.float32)
156
- luma = roi.mean(axis=2)
157
- saturation = roi.max(axis=2) - roi.min(axis=2)
158
- grayish = saturation < DOUBAO_MAX_SATURATION
159
- sigma = max(4.0, height * 0.4)
160
- local_bg = cv2.GaussianBlur(luma, (0, 0), sigmaX=sigma, sigmaY=sigma)
161
- tophat = luma - local_bg
162
- glyph = (
163
- grayish
164
- & (tophat > DOUBAO_TOPHAT_DELTA)
165
- & (luma > DOUBAO_LOGO_MIN_LUMA)
166
- ).astype(np.uint8) * 255
167
- glyph = cv2.morphologyEx(glyph, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
168
- return cv2.morphologyEx(glyph, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
169
-
170
-
171
- def _doubao_template_match_score(box_mask: np.ndarray, image_width: int) -> float:
172
- silhouette = _load_doubao_silhouette()
173
- if box_mask.size == 0:
174
- return 0.0
175
- gw = min(box_mask.shape[1] - 1, max(8, int(DOUBAO_ALPHA_WIDTH_FRAC * image_width)))
176
- gh = min(box_mask.shape[0] - 1, max(4, int(DOUBAO_ALPHA_HEIGHT_FRAC * image_width)))
177
- if gw < 8 or gh < 4:
178
- return 0.0
179
- template = cv2.resize(silhouette, (gw, gh), interpolation=cv2.INTER_NEAREST)
180
- return float(cv2.matchTemplate(box_mask, template, cv2.TM_CCOEFF_NORMED).max())
181
-
182
-
183
- def _doubao_aligned_alpha(image_bgr: np.ndarray) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
184
- alpha = _load_doubao_alpha()
185
- silhouette = _load_doubao_silhouette()
186
- image_width = image_bgr.shape[1]
187
- bx, by, bw, bh = _doubao_locate(image_bgr)
188
- box_mask = _doubao_extract_candidate_mask(image_bgr, (bx, by, bw, bh))
189
- expected = DOUBAO_ALPHA_WIDTH_FRAC * image_width
190
- best: tuple[float, int, int, int, int] | None = None
191
- for scale in np.linspace(*DOUBAO_ALPHA_ALIGN_SEARCH):
192
- gw = int(expected * scale)
193
- gh = int(DOUBAO_ALPHA_HEIGHT_FRAC * image_width * scale)
194
- if gw < 8 or gh < 4 or gw >= bw or gh >= bh:
195
- continue
196
- template = cv2.resize(silhouette, (gw, gh), interpolation=cv2.INTER_NEAREST)
197
- _, score, _, top_left = cv2.minMaxLoc(cv2.matchTemplate(box_mask, template, cv2.TM_CCOEFF_NORMED))
198
- if best is None or score > best[0]:
199
- best = (score, gw, gh, top_left[0], top_left[1])
200
- if best is None:
201
- return None
202
- _, gw, gh, ox, oy = best
203
- return cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_LINEAR), (bx + ox, by + oy, gw, gh)
204
-
205
-
206
- def _doubao_fixed_alpha(image_bgr: np.ndarray) -> tuple[np.ndarray, tuple[int, int, int, int]]:
207
- alpha = _load_doubao_alpha()
208
- image_h, image_w = image_bgr.shape[:2]
209
- gw = min(image_w, max(1, int(DOUBAO_ALPHA_WIDTH_FRAC * image_w)))
210
- gh = min(image_h, max(1, int(DOUBAO_ALPHA_HEIGHT_FRAC * image_w)))
211
- ax = max(0, image_w - int(DOUBAO_ALPHA_MARGIN_RIGHT_FRAC * image_w) - gw)
212
- ay = max(0, image_h - int(DOUBAO_ALPHA_MARGIN_BOTTOM_FRAC * image_w) - gh)
213
- return cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_LINEAR), (ax, ay, gw, gh)
214
-
215
-
216
- def generate_doubao_template_bbox_mask(image_bgr: np.ndarray) -> tuple[np.ndarray, dict[str, object]]:
217
- if min(image_bgr.shape[:2]) < 200:
218
- raise ValueError("Image is too small for Doubao watermark detection.")
219
- loc = _doubao_locate(image_bgr)
220
- candidate_mask = _doubao_extract_candidate_mask(image_bgr, loc)
221
- coverage = float((candidate_mask > 0).sum()) / float(max(1, loc[2] * loc[3]))
222
- confidence = _doubao_template_match_score(candidate_mask, image_bgr.shape[1]) if coverage >= DOUBAO_DETECT_MIN_COVERAGE else 0.0
223
- detected = confidence >= DOUBAO_DETECT_NCC_THRESHOLD
224
- if not detected:
225
- raise ValueError(f"Doubao watermark not detected. confidence={confidence:.3f} coverage={coverage:.3f}")
226
-
227
- placed = _doubao_aligned_alpha(image_bgr) or _doubao_fixed_alpha(image_bgr)
228
- alpha_block, (x, y, width, height) = placed
229
- silhouette = (alpha_block > DOUBAO_RESIDUAL_ALPHA_FLOOR).astype(np.uint8) * 255
230
- if int((silhouette > 0).sum()) == 0:
231
- raise ValueError("Detected Doubao watermark but generated an empty mask.")
232
- glyph_mask = np.zeros(image_bgr.shape[:2], np.uint8)
233
- glyph_mask[y : y + height, x : x + width] = silhouette
234
- if DOUBAO_TEMPLATE_DILATE > 0:
235
- kernel = cv2.getStructuringElement(
236
- cv2.MORPH_ELLIPSE,
237
- (2 * DOUBAO_TEMPLATE_DILATE + 1, 2 * DOUBAO_TEMPLATE_DILATE + 1),
238
- )
239
- glyph_mask = cv2.dilate(glyph_mask, kernel)
240
-
241
- ys, xs = np.where(glyph_mask > 0)
242
- if len(xs) == 0 or len(ys) == 0:
243
- raise ValueError("Detected Doubao watermark but bbox mask is empty.")
244
- bbox_mask = np.zeros_like(glyph_mask)
245
- x0, x1 = int(xs.min()), int(xs.max()) + 1
246
- y0, y1 = int(ys.min()), int(ys.max()) + 1
247
- bbox_mask[y0:y1, x0:x1] = 255
248
- return bbox_mask, {
249
- "confidence": round(confidence, 4),
250
- "coverage": round(coverage, 4),
251
- "bbox": [x0, y0, x1 - x0, y1 - y0],
252
- }
253
-
254
-
255
- def _mask_bbox(mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
256
- ys, xs = np.where(mask > 0)
257
- if len(xs) == 0 or len(ys) == 0:
258
- raise ValueError("Mask is empty.")
259
- return ys, xs
260
-
261
-
262
- def _ceil_modulo(value: int, mod: int) -> int:
263
- if value % mod == 0:
264
- return value
265
- return (value // mod + 1) * mod
266
-
267
-
268
- def _pad_img_to_modulo(
269
- image: np.ndarray,
270
- *,
271
- mod: int,
272
- square: bool = False,
273
- min_size: int | None = None,
274
- ) -> np.ndarray:
275
- if image.ndim == 2:
276
- image = image[:, :, None]
277
- height, width = image.shape[:2]
278
- out_height = _ceil_modulo(height, mod)
279
- out_width = _ceil_modulo(width, mod)
280
- if min_size is not None:
281
- out_width = max(min_size, out_width)
282
- out_height = max(min_size, out_height)
283
- if square:
284
- side = max(out_height, out_width)
285
- out_height = side
286
- out_width = side
287
- return np.pad(
288
- image,
289
- ((0, out_height - height), (0, out_width - width), (0, 0)),
290
- mode="symmetric",
291
- )
292
-
293
-
294
- def _resize_max_size(image: np.ndarray, size_limit: int, interpolation: int = cv2.INTER_CUBIC) -> np.ndarray:
295
- height, width = image.shape[:2]
296
- if max(height, width) <= size_limit:
297
- return image
298
- ratio = size_limit / max(height, width)
299
- return cv2.resize(image, (int(width * ratio + 0.5), int(height * ratio + 0.5)), interpolation=interpolation)
300
-
301
-
302
- def _norm_img(image: np.ndarray) -> np.ndarray:
303
- if image.ndim == 2:
304
- image = image[:, :, None]
305
- image = np.transpose(image, (2, 0, 1))
306
- return image.astype("float32") / 255.0
307
-
308
-
309
- def _boxes_from_mask(mask: np.ndarray) -> list[np.ndarray]:
310
- _, thresh = cv2.threshold(mask, 127, 255, 0)
311
- contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
312
- boxes: list[np.ndarray] = []
313
- height, width = mask.shape[:2]
314
- for contour in contours:
315
- x, y, box_width, box_height = cv2.boundingRect(contour)
316
- box = np.array([x, y, x + box_width, y + box_height]).astype(int)
317
- box[::2] = np.clip(box[::2], 0, width)
318
- box[1::2] = np.clip(box[1::2], 0, height)
319
- boxes.append(box)
320
- return boxes
321
-
322
-
323
- def _crop_box(
324
- image: np.ndarray,
325
- mask: np.ndarray,
326
- box: np.ndarray,
327
- *,
328
- margin: int,
329
- ) -> tuple[np.ndarray, np.ndarray, tuple[int, int, int, int]]:
330
- box_h = int(box[3] - box[1])
331
- box_w = int(box[2] - box[0])
332
- cx = int((box[0] + box[2]) // 2)
333
- cy = int((box[1] + box[3]) // 2)
334
- img_h, img_w = image.shape[:2]
335
-
336
- crop_w = box_w + margin * 2
337
- crop_h = box_h + margin * 2
338
- raw_l = cx - crop_w // 2
339
- raw_r = cx + crop_w // 2
340
- raw_t = cy - crop_h // 2
341
- raw_b = cy + crop_h // 2
342
-
343
- left = max(raw_l, 0)
344
- right = min(raw_r, img_w)
345
- top = max(raw_t, 0)
346
- bottom = min(raw_b, img_h)
347
-
348
- if raw_l < 0:
349
- right += abs(raw_l)
350
- if raw_r > img_w:
351
- left -= raw_r - img_w
352
- if raw_t < 0:
353
- bottom += abs(raw_t)
354
- if raw_b > img_h:
355
- top -= raw_b - img_h
356
-
357
- left = max(left, 0)
358
- right = min(right, img_w)
359
- top = max(top, 0)
360
- bottom = min(bottom, img_h)
361
- return image[top:bottom, left:right, :], mask[top:bottom, left:right], (left, top, right, bottom)
362
-
363
-
364
- def _erase_lama_onnx(image_bgr: np.ndarray, mask: np.ndarray) -> np.ndarray:
365
- session = _load_lama_session()
366
- inputs = session.get_inputs()
367
- image_name = inputs[0].name
368
- mask_name = inputs[1].name
369
- dims = inputs[0].shape
370
- size = next((dim for dim in reversed(dims) if isinstance(dim, int) and dim > 1), 512)
371
-
372
- height, width = image_bgr.shape[:2]
373
- ys, xs = _mask_bbox(mask)
374
-
375
- pad = max(16, int(0.4 * max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1)))
376
- cx0 = max(0, int(xs.min()) - pad)
377
- cy0 = max(0, int(ys.min()) - pad)
378
- cx1 = min(width, int(xs.max()) + 1 + pad)
379
- cy1 = min(height, int(ys.max()) + 1 + pad)
380
-
381
- crop = image_bgr[cy0:cy1, cx0:cx1]
382
- crop_mask = mask[cy0:cy1, cx0:cx1]
383
- crop_h, crop_w = crop.shape[:2]
384
-
385
- crop_rs = cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA)
386
- mask_rs = cv2.resize(crop_mask, (size, size), interpolation=cv2.INTER_NEAREST)
387
-
388
- image_input = cv2.cvtColor(crop_rs, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
389
- image_input = np.transpose(image_input, (2, 0, 1))[None]
390
- mask_input = (mask_rs > 127).astype(np.float32)[None, None]
391
-
392
- output = session.run(None, {image_name: image_input, mask_name: mask_input})[0]
393
- output = np.asarray(output)[0]
394
- output = np.transpose(output, (1, 2, 0))
395
- if float(output.max()) <= 1.5:
396
- output = output * 255.0
397
- output = np.clip(output, 0, 255).astype(np.uint8)
398
- output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
399
- output_crop = cv2.resize(output_bgr, (crop_w, crop_h), interpolation=cv2.INTER_LINEAR)
400
-
401
- result = image_bgr.copy()
402
- region = result[cy0:cy1, cx0:cx1]
403
- paste = crop_mask > 127
404
- region[paste] = output_crop[paste]
405
- result[cy0:cy1, cx0:cx1] = region
406
- return result
407
-
408
-
409
- @torch.inference_mode()
410
- def _migan_forward(model: torch.jit.ScriptModule, image_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
411
- image = _norm_img(image_rgb) * 2 - 1
412
- mask_binary = (mask > 120).astype(np.uint8) * 255
413
- mask_input = _norm_img(mask_binary)
414
- image_tensor = torch.from_numpy(image).unsqueeze(0)
415
- mask_tensor = torch.from_numpy(mask_input).unsqueeze(0)
416
- erased = image_tensor * (1 - mask_tensor)
417
- model_input = torch.cat([0.5 - mask_tensor, erased], dim=1)
418
- output = model(model_input)
419
- output = (
420
- (output.permute(0, 2, 3, 1) * 127.5 + 127.5)
421
- .round()
422
- .clamp(0, 255)
423
- .to(torch.uint8)
424
- )
425
- result = output[0].cpu().numpy()
426
- return cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
427
-
428
-
429
- def _migan_pad_forward(model: torch.jit.ScriptModule, image_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
430
- origin_h, origin_w = image_rgb.shape[:2]
431
- padded_image = _pad_img_to_modulo(image_rgb, mod=512, square=True, min_size=512)
432
- padded_mask = _pad_img_to_modulo(mask, mod=512, square=True, min_size=512)
433
- return _migan_forward(model, padded_image, padded_mask)[:origin_h, :origin_w, :]
434
-
435
-
436
- def _erase_migan_torchscript(image_bgr: np.ndarray, mask: np.ndarray) -> np.ndarray:
437
- model = _load_migan_model()
438
- image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
439
- if image_rgb.shape[:2] == (512, 512):
440
- return _migan_pad_forward(model, image_rgb, mask)
441
-
442
- result = image_rgb[:, :, ::-1].copy()
443
- for box in _boxes_from_mask(mask):
444
- crop_img, crop_mask, (left, top, right, bottom) = _crop_box(image_rgb, mask, box, margin=128)
445
- origin_size = crop_img.shape[:2]
446
- resized_img = _resize_max_size(crop_img, 512)
447
- resized_mask = _resize_max_size(crop_mask, 512)
448
- inpaint = _migan_pad_forward(model, resized_img, resized_mask)
449
- inpaint = cv2.resize(inpaint, (origin_size[1], origin_size[0]), interpolation=cv2.INTER_CUBIC)
450
- keep = crop_mask < 127
451
- inpaint[keep] = crop_img[:, :, ::-1][keep]
452
- result[top:bottom, left:right, :] = inpaint
453
- return result
454
-
455
-
456
- def _normalize_mode(mode: str | None) -> str:
457
- requested = (mode or DEFAULT_MODE).strip().lower()
458
- aliases = {
459
- "lama": "quality",
460
- "high": "quality",
461
- "high-quality": "quality",
462
- "migan": "fast",
463
- "quick": "fast",
464
- }
465
- normalized = aliases.get(requested, requested)
466
- if normalized not in MODES:
467
- raise ValueError(f"Unsupported mode: {mode}. Supported modes: {', '.join(MODES)}")
468
- return normalized
469
-
470
-
471
- def run_inpaint(image: Image.Image, mask: Image.Image, *, mode: str | None = None) -> tuple[Image.Image, str, float]:
472
- normalized_mode = _normalize_mode(mode)
473
- rgb_image = np.array(image.convert("RGB"), copy=True)
474
- image_bgr = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
475
- prepared_mask = _prepare_mask(mask, image.size)
476
-
477
- start = time.time()
478
- with MODEL_LOCK:
479
- if normalized_mode == "quality":
480
- result_bgr = _erase_lama_onnx(image_bgr, prepared_mask)
481
- elif normalized_mode == "fast":
482
- result_bgr = _erase_migan_torchscript(image_bgr, prepared_mask)
483
- else:
484
- raise ValueError(f"Unsupported mode: {normalized_mode}")
485
- elapsed = time.time() - start
486
-
487
- result_rgb = cv2.cvtColor(result_bgr.astype(np.uint8), cv2.COLOR_BGR2RGB)
488
- return Image.fromarray(result_rgb), normalized_mode, elapsed
489
-
490
-
491
- def run_remove_doubao(image: Image.Image, *, mode: str | None = None) -> tuple[Image.Image, str, float, dict[str, object]]:
492
- rgb_image = np.array(image.convert("RGB"), copy=True)
493
- image_bgr = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
494
- mask, metadata = generate_doubao_template_bbox_mask(image_bgr)
495
- mask_image = Image.fromarray(mask, mode="L")
496
- result, normalized_mode, elapsed = run_inpaint(image, mask_image, mode=mode)
497
- return result, normalized_mode, elapsed, metadata
498
-
499
-
500
- @app.on_event("startup")
501
- def startup() -> None:
502
- if not LOAD_MODEL_ON_STARTUP:
503
- return
504
- if DEFAULT_MODE == "fast":
505
- _load_migan_model()
506
- else:
507
- _load_lama_session()
508
-
509
-
510
- @app.get("/")
511
- def root() -> dict[str, object]:
512
- return health()
513
-
514
-
515
- @app.get("/health")
516
- def health() -> dict[str, object]:
517
- return {
518
- "status": "ok",
519
- "mode": "dual-selected-backends",
520
- "default_mode": DEFAULT_MODE,
521
- "supported_modes": MODES,
522
- "loaded": {
523
- "quality": LAMA_SESSION is not None,
524
- "fast": MIGAN_MODEL is not None,
525
- },
526
- }
527
-
528
-
529
- @app.post("/inpaint")
530
- async def inpaint(
531
- image: UploadFile = File(...),
532
- mask: UploadFile = File(...),
533
- mode: str = Form(DEFAULT_MODE),
534
- ) -> Response:
535
- try:
536
- image_bytes = await image.read()
537
- mask_bytes = await mask.read()
538
- pil_image = Image.open(io.BytesIO(image_bytes))
539
- pil_mask = Image.open(io.BytesIO(mask_bytes))
540
- result, normalized_mode, elapsed = run_inpaint(pil_image, pil_mask, mode=mode)
541
- output = io.BytesIO()
542
- result.save(output, format="PNG")
543
- return Response(
544
- content=output.getvalue(),
545
- media_type="image/png",
546
- headers={
547
- "X-Inpaint-Mode": normalized_mode,
548
- "X-Inpaint-Backend": MODES[normalized_mode],
549
- "X-Inpaint-Elapsed": f"{elapsed:.3f}",
550
- },
551
- )
552
- except Exception as exc:
553
- raise HTTPException(status_code=500, detail=str(exc)) from exc
554
-
555
-
556
- @app.post("/remove-doubao")
557
- async def remove_doubao(
558
- image: UploadFile = File(...),
559
- mode: str = Form(DEFAULT_MODE),
560
- ) -> Response:
561
- try:
562
- image_bytes = await image.read()
563
- pil_image = Image.open(io.BytesIO(image_bytes))
564
- result, normalized_mode, elapsed, metadata = run_remove_doubao(pil_image, mode=mode)
565
- output = io.BytesIO()
566
- result.save(output, format="PNG")
567
- return Response(
568
- content=output.getvalue(),
569
- media_type="image/png",
570
- headers={
571
- "X-Inpaint-Mode": normalized_mode,
572
- "X-Inpaint-Backend": MODES[normalized_mode],
573
- "X-Inpaint-Elapsed": f"{elapsed:.3f}",
574
- "X-Doubao-Confidence": str(metadata["confidence"]),
575
- "X-Doubao-Coverage": str(metadata["coverage"]),
576
- "X-Doubao-BBox": ",".join(str(v) for v in metadata["bbox"]),
577
- },
578
- )
579
- except Exception as exc:
580
- raise HTTPException(status_code=500, detail=str(exc)) from exc
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import hashlib
4
+ import io
5
+ import os
6
+ import threading
7
+ import time
8
+ from pathlib import Path
9
+ from urllib.parse import urlparse
10
+
11
+ import cv2
12
+ import numpy as np
13
+ import onnxruntime as ort
14
+ from fastapi import FastAPI, File, Form, HTTPException, UploadFile
15
+ from fastapi.responses import Response
16
+ from huggingface_hub import hf_hub_download
17
+ from PIL import Image
18
+
19
+
20
+ LAMA_MODEL_REPO = os.getenv("LAMA_MODEL_REPO", "Carve/LaMa-ONNX")
21
+ LAMA_MODEL_FILE = os.getenv("LAMA_MODEL_FILE", "lama_fp32.onnx")
22
+ MIGAN_MODEL_URL = os.getenv(
23
+ "MIGAN_MODEL_URL",
24
+ "https://github.com/Sanster/models/releases/download/migan/migan_traced.pt",
25
+ )
26
+ MIGAN_MODEL_MD5 = os.getenv("MIGAN_MODEL_MD5", "76eb3b1a71c400ee3290524f7a11b89c")
27
+ DEFAULT_MODE = os.getenv("DEFAULT_MODE", "quality")
28
+ LOAD_MODEL_ON_STARTUP = os.getenv("LOAD_MODEL_ON_STARTUP", "1") == "1"
29
+ PROVIDERS = ["CPUExecutionProvider"]
30
+ MODES = {
31
+ "quality": "selfbuilt_lama_onnx",
32
+ "fast": "iopaint_migan_torchscript",
33
+ }
34
+
35
+ DOUBAO_ALPHA_ASSET = Path(__file__).with_name("doubao_alpha.png")
36
+ DOUBAO_WIDTH_FRAC = 0.22
37
+ DOUBAO_HEIGHT_FRAC = 0.075
38
+ DOUBAO_MARGIN_RIGHT_FRAC = 0.004
39
+ DOUBAO_MARGIN_BOTTOM_FRAC = 0.004
40
+ DOUBAO_MAX_SATURATION = 55
41
+ DOUBAO_LOGO_MIN_LUMA = 150
42
+ DOUBAO_TOPHAT_DELTA = 12
43
+ DOUBAO_DETECT_MIN_COVERAGE = 0.04
44
+ DOUBAO_DETECT_NCC_THRESHOLD = 0.4
45
+ DOUBAO_ALPHA_WIDTH_FRAC = 0.1636
46
+ DOUBAO_ALPHA_HEIGHT_FRAC = 0.0405
47
+ DOUBAO_ALPHA_MARGIN_RIGHT_FRAC = 0.0132
48
+ DOUBAO_ALPHA_MARGIN_BOTTOM_FRAC = 0.0166
49
+ DOUBAO_ALPHA_ALIGN_SEARCH = (0.88, 1.12, 25)
50
+ DOUBAO_RESIDUAL_ALPHA_FLOOR = 0.05
51
+ DOUBAO_TEMPLATE_DILATE = 4
52
+
53
+ LAMA_SESSION: ort.InferenceSession | None = None
54
+ MIGAN_MODEL: object | None = None
55
+ DOUBAO_ALPHA: np.ndarray | None = None
56
+ DOUBAO_SILHOUETTE: np.ndarray | None = None
57
+ MODEL_LOCK = threading.RLock()
58
+
59
+ app = FastAPI(title="Image Services Inpaint")
60
+
61
+
62
+ def _md5sum(path: Path) -> str:
63
+ digest = hashlib.md5()
64
+ with path.open("rb") as handle:
65
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
66
+ digest.update(chunk)
67
+ return digest.hexdigest()
68
+
69
+
70
+ def _cache_path_by_url(url: str) -> Path:
71
+ parts = urlparse(url)
72
+ model_dir = Path(os.getenv("TORCH_HOME", "/tmp/torch")) / "checkpoints"
73
+ model_dir.mkdir(parents=True, exist_ok=True)
74
+ return model_dir / Path(parts.path).name
75
+
76
+
77
+ def _download_torchscript(url: str, expected_md5: str) -> Path:
78
+ from torch.hub import download_url_to_file
79
+
80
+ path = _cache_path_by_url(url)
81
+ if not path.exists():
82
+ download_url_to_file(url, str(path), None, progress=True)
83
+ actual_md5 = _md5sum(path)
84
+ if actual_md5 != expected_md5:
85
+ try:
86
+ path.unlink()
87
+ finally:
88
+ raise RuntimeError(f"Model md5 mismatch: {actual_md5} != {expected_md5}")
89
+ return path
90
+
91
+
92
+ def _load_lama_session() -> ort.InferenceSession:
93
+ global LAMA_SESSION
94
+ if LAMA_SESSION is not None:
95
+ return LAMA_SESSION
96
+ with MODEL_LOCK:
97
+ if LAMA_SESSION is not None:
98
+ return LAMA_SESSION
99
+ model_path = hf_hub_download(repo_id=LAMA_MODEL_REPO, filename=LAMA_MODEL_FILE)
100
+ LAMA_SESSION = ort.InferenceSession(model_path, providers=PROVIDERS)
101
+ return LAMA_SESSION
102
+
103
+
104
+ def _load_migan_model() -> object:
105
+ global MIGAN_MODEL
106
+ if MIGAN_MODEL is not None:
107
+ return MIGAN_MODEL
108
+ with MODEL_LOCK:
109
+ if MIGAN_MODEL is not None:
110
+ return MIGAN_MODEL
111
+ model_path = _download_torchscript(MIGAN_MODEL_URL, MIGAN_MODEL_MD5)
112
+ import torch
113
+
114
+ MIGAN_MODEL = torch.jit.load(str(model_path), map_location="cpu").eval()
115
+ return MIGAN_MODEL
116
+
117
+
118
+ def _prepare_mask(mask: Image.Image, size: tuple[int, int]) -> np.ndarray:
119
+ arr = np.array(mask.convert("L").resize(size, Image.Resampling.NEAREST), copy=True)
120
+ return np.where(arr > 127, 255, 0).astype(np.uint8)
121
+
122
+
123
+ def _load_doubao_alpha() -> np.ndarray:
124
+ global DOUBAO_ALPHA
125
+ if DOUBAO_ALPHA is not None:
126
+ return DOUBAO_ALPHA
127
+ alpha = cv2.imread(str(DOUBAO_ALPHA_ASSET), cv2.IMREAD_GRAYSCALE)
128
+ if alpha is None:
129
+ raise RuntimeError(f"Missing Doubao alpha asset: {DOUBAO_ALPHA_ASSET}")
130
+ DOUBAO_ALPHA = alpha.astype(np.float32) / 255.0
131
+ return DOUBAO_ALPHA
132
+
133
+
134
+ def _load_doubao_silhouette() -> np.ndarray:
135
+ global DOUBAO_SILHOUETTE
136
+ if DOUBAO_SILHOUETTE is not None:
137
+ return DOUBAO_SILHOUETTE
138
+ DOUBAO_SILHOUETTE = (_load_doubao_alpha() > 0.15).astype(np.uint8) * 255
139
+ return DOUBAO_SILHOUETTE
140
+
141
+
142
+ def _doubao_locate(image_bgr: np.ndarray) -> tuple[int, int, int, int]:
143
+ height, width = image_bgr.shape[:2]
144
+ mark_w = max(40, int(width * DOUBAO_WIDTH_FRAC))
145
+ mark_h = max(16, int(width * DOUBAO_HEIGHT_FRAC))
146
+ margin_right = max(4, int(width * DOUBAO_MARGIN_RIGHT_FRAC))
147
+ margin_bottom = max(4, int(width * DOUBAO_MARGIN_BOTTOM_FRAC))
148
+ x = max(0, width - margin_right - mark_w)
149
+ y = max(0, height - margin_bottom - mark_h)
150
+ return x, y, min(mark_w, width - x), min(mark_h, height - y)
151
+
152
+
153
+ def _doubao_extract_candidate_mask(image_bgr: np.ndarray, loc: tuple[int, int, int, int]) -> np.ndarray:
154
+ x, y, width, height = loc
155
+ if width < 16 or height < 16:
156
+ return np.zeros((height, width), np.uint8)
157
+ roi = image_bgr[y : y + height, x : x + width].astype(np.float32)
158
+ luma = roi.mean(axis=2)
159
+ saturation = roi.max(axis=2) - roi.min(axis=2)
160
+ grayish = saturation < DOUBAO_MAX_SATURATION
161
+ sigma = max(4.0, height * 0.4)
162
+ local_bg = cv2.GaussianBlur(luma, (0, 0), sigmaX=sigma, sigmaY=sigma)
163
+ tophat = luma - local_bg
164
+ glyph = (
165
+ grayish
166
+ & (tophat > DOUBAO_TOPHAT_DELTA)
167
+ & (luma > DOUBAO_LOGO_MIN_LUMA)
168
+ ).astype(np.uint8) * 255
169
+ glyph = cv2.morphologyEx(glyph, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
170
+ return cv2.morphologyEx(glyph, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
171
+
172
+
173
+ def _doubao_template_match_score(box_mask: np.ndarray, image_width: int) -> float:
174
+ silhouette = _load_doubao_silhouette()
175
+ if box_mask.size == 0:
176
+ return 0.0
177
+ gw = min(box_mask.shape[1] - 1, max(8, int(DOUBAO_ALPHA_WIDTH_FRAC * image_width)))
178
+ gh = min(box_mask.shape[0] - 1, max(4, int(DOUBAO_ALPHA_HEIGHT_FRAC * image_width)))
179
+ if gw < 8 or gh < 4:
180
+ return 0.0
181
+ template = cv2.resize(silhouette, (gw, gh), interpolation=cv2.INTER_NEAREST)
182
+ return float(cv2.matchTemplate(box_mask, template, cv2.TM_CCOEFF_NORMED).max())
183
+
184
+
185
+ def _doubao_aligned_alpha(image_bgr: np.ndarray) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
186
+ alpha = _load_doubao_alpha()
187
+ silhouette = _load_doubao_silhouette()
188
+ image_width = image_bgr.shape[1]
189
+ bx, by, bw, bh = _doubao_locate(image_bgr)
190
+ box_mask = _doubao_extract_candidate_mask(image_bgr, (bx, by, bw, bh))
191
+ expected = DOUBAO_ALPHA_WIDTH_FRAC * image_width
192
+ best: tuple[float, int, int, int, int] | None = None
193
+ for scale in np.linspace(*DOUBAO_ALPHA_ALIGN_SEARCH):
194
+ gw = int(expected * scale)
195
+ gh = int(DOUBAO_ALPHA_HEIGHT_FRAC * image_width * scale)
196
+ if gw < 8 or gh < 4 or gw >= bw or gh >= bh:
197
+ continue
198
+ template = cv2.resize(silhouette, (gw, gh), interpolation=cv2.INTER_NEAREST)
199
+ _, score, _, top_left = cv2.minMaxLoc(cv2.matchTemplate(box_mask, template, cv2.TM_CCOEFF_NORMED))
200
+ if best is None or score > best[0]:
201
+ best = (score, gw, gh, top_left[0], top_left[1])
202
+ if best is None:
203
+ return None
204
+ _, gw, gh, ox, oy = best
205
+ return cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_LINEAR), (bx + ox, by + oy, gw, gh)
206
+
207
+
208
+ def _doubao_fixed_alpha(image_bgr: np.ndarray) -> tuple[np.ndarray, tuple[int, int, int, int]]:
209
+ alpha = _load_doubao_alpha()
210
+ image_h, image_w = image_bgr.shape[:2]
211
+ gw = min(image_w, max(1, int(DOUBAO_ALPHA_WIDTH_FRAC * image_w)))
212
+ gh = min(image_h, max(1, int(DOUBAO_ALPHA_HEIGHT_FRAC * image_w)))
213
+ ax = max(0, image_w - int(DOUBAO_ALPHA_MARGIN_RIGHT_FRAC * image_w) - gw)
214
+ ay = max(0, image_h - int(DOUBAO_ALPHA_MARGIN_BOTTOM_FRAC * image_w) - gh)
215
+ return cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_LINEAR), (ax, ay, gw, gh)
216
+
217
+
218
+ def generate_doubao_template_bbox_mask(image_bgr: np.ndarray) -> tuple[np.ndarray, dict[str, object]]:
219
+ if min(image_bgr.shape[:2]) < 200:
220
+ raise ValueError("Image is too small for Doubao watermark detection.")
221
+ loc = _doubao_locate(image_bgr)
222
+ candidate_mask = _doubao_extract_candidate_mask(image_bgr, loc)
223
+ coverage = float((candidate_mask > 0).sum()) / float(max(1, loc[2] * loc[3]))
224
+ confidence = _doubao_template_match_score(candidate_mask, image_bgr.shape[1]) if coverage >= DOUBAO_DETECT_MIN_COVERAGE else 0.0
225
+ detected = confidence >= DOUBAO_DETECT_NCC_THRESHOLD
226
+ if not detected:
227
+ raise ValueError(f"Doubao watermark not detected. confidence={confidence:.3f} coverage={coverage:.3f}")
228
+
229
+ placed = _doubao_aligned_alpha(image_bgr) or _doubao_fixed_alpha(image_bgr)
230
+ alpha_block, (x, y, width, height) = placed
231
+ silhouette = (alpha_block > DOUBAO_RESIDUAL_ALPHA_FLOOR).astype(np.uint8) * 255
232
+ if int((silhouette > 0).sum()) == 0:
233
+ raise ValueError("Detected Doubao watermark but generated an empty mask.")
234
+ glyph_mask = np.zeros(image_bgr.shape[:2], np.uint8)
235
+ glyph_mask[y : y + height, x : x + width] = silhouette
236
+ if DOUBAO_TEMPLATE_DILATE > 0:
237
+ kernel = cv2.getStructuringElement(
238
+ cv2.MORPH_ELLIPSE,
239
+ (2 * DOUBAO_TEMPLATE_DILATE + 1, 2 * DOUBAO_TEMPLATE_DILATE + 1),
240
+ )
241
+ glyph_mask = cv2.dilate(glyph_mask, kernel)
242
+
243
+ ys, xs = np.where(glyph_mask > 0)
244
+ if len(xs) == 0 or len(ys) == 0:
245
+ raise ValueError("Detected Doubao watermark but bbox mask is empty.")
246
+ bbox_mask = np.zeros_like(glyph_mask)
247
+ x0, x1 = int(xs.min()), int(xs.max()) + 1
248
+ y0, y1 = int(ys.min()), int(ys.max()) + 1
249
+ bbox_mask[y0:y1, x0:x1] = 255
250
+ return bbox_mask, {
251
+ "confidence": round(confidence, 4),
252
+ "coverage": round(coverage, 4),
253
+ "bbox": [x0, y0, x1 - x0, y1 - y0],
254
+ }
255
+
256
+
257
+ def _mask_bbox(mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
258
+ ys, xs = np.where(mask > 0)
259
+ if len(xs) == 0 or len(ys) == 0:
260
+ raise ValueError("Mask is empty.")
261
+ return ys, xs
262
+
263
+
264
+ def _ceil_modulo(value: int, mod: int) -> int:
265
+ if value % mod == 0:
266
+ return value
267
+ return (value // mod + 1) * mod
268
+
269
+
270
+ def _pad_img_to_modulo(
271
+ image: np.ndarray,
272
+ *,
273
+ mod: int,
274
+ square: bool = False,
275
+ min_size: int | None = None,
276
+ ) -> np.ndarray:
277
+ if image.ndim == 2:
278
+ image = image[:, :, None]
279
+ height, width = image.shape[:2]
280
+ out_height = _ceil_modulo(height, mod)
281
+ out_width = _ceil_modulo(width, mod)
282
+ if min_size is not None:
283
+ out_width = max(min_size, out_width)
284
+ out_height = max(min_size, out_height)
285
+ if square:
286
+ side = max(out_height, out_width)
287
+ out_height = side
288
+ out_width = side
289
+ return np.pad(
290
+ image,
291
+ ((0, out_height - height), (0, out_width - width), (0, 0)),
292
+ mode="symmetric",
293
+ )
294
+
295
+
296
+ def _resize_max_size(image: np.ndarray, size_limit: int, interpolation: int = cv2.INTER_CUBIC) -> np.ndarray:
297
+ height, width = image.shape[:2]
298
+ if max(height, width) <= size_limit:
299
+ return image
300
+ ratio = size_limit / max(height, width)
301
+ return cv2.resize(image, (int(width * ratio + 0.5), int(height * ratio + 0.5)), interpolation=interpolation)
302
+
303
+
304
+ def _norm_img(image: np.ndarray) -> np.ndarray:
305
+ if image.ndim == 2:
306
+ image = image[:, :, None]
307
+ image = np.transpose(image, (2, 0, 1))
308
+ return image.astype("float32") / 255.0
309
+
310
+
311
+ def _boxes_from_mask(mask: np.ndarray) -> list[np.ndarray]:
312
+ _, thresh = cv2.threshold(mask, 127, 255, 0)
313
+ contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
314
+ boxes: list[np.ndarray] = []
315
+ height, width = mask.shape[:2]
316
+ for contour in contours:
317
+ x, y, box_width, box_height = cv2.boundingRect(contour)
318
+ box = np.array([x, y, x + box_width, y + box_height]).astype(int)
319
+ box[::2] = np.clip(box[::2], 0, width)
320
+ box[1::2] = np.clip(box[1::2], 0, height)
321
+ boxes.append(box)
322
+ return boxes
323
+
324
+
325
+ def _crop_box(
326
+ image: np.ndarray,
327
+ mask: np.ndarray,
328
+ box: np.ndarray,
329
+ *,
330
+ margin: int,
331
+ ) -> tuple[np.ndarray, np.ndarray, tuple[int, int, int, int]]:
332
+ box_h = int(box[3] - box[1])
333
+ box_w = int(box[2] - box[0])
334
+ cx = int((box[0] + box[2]) // 2)
335
+ cy = int((box[1] + box[3]) // 2)
336
+ img_h, img_w = image.shape[:2]
337
+
338
+ crop_w = box_w + margin * 2
339
+ crop_h = box_h + margin * 2
340
+ raw_l = cx - crop_w // 2
341
+ raw_r = cx + crop_w // 2
342
+ raw_t = cy - crop_h // 2
343
+ raw_b = cy + crop_h // 2
344
+
345
+ left = max(raw_l, 0)
346
+ right = min(raw_r, img_w)
347
+ top = max(raw_t, 0)
348
+ bottom = min(raw_b, img_h)
349
+
350
+ if raw_l < 0:
351
+ right += abs(raw_l)
352
+ if raw_r > img_w:
353
+ left -= raw_r - img_w
354
+ if raw_t < 0:
355
+ bottom += abs(raw_t)
356
+ if raw_b > img_h:
357
+ top -= raw_b - img_h
358
+
359
+ left = max(left, 0)
360
+ right = min(right, img_w)
361
+ top = max(top, 0)
362
+ bottom = min(bottom, img_h)
363
+ return image[top:bottom, left:right, :], mask[top:bottom, left:right], (left, top, right, bottom)
364
+
365
+
366
+ def _erase_lama_onnx(image_bgr: np.ndarray, mask: np.ndarray) -> np.ndarray:
367
+ session = _load_lama_session()
368
+ inputs = session.get_inputs()
369
+ image_name = inputs[0].name
370
+ mask_name = inputs[1].name
371
+ dims = inputs[0].shape
372
+ size = next((dim for dim in reversed(dims) if isinstance(dim, int) and dim > 1), 512)
373
+
374
+ height, width = image_bgr.shape[:2]
375
+ ys, xs = _mask_bbox(mask)
376
+
377
+ pad = max(16, int(0.4 * max(xs.max() - xs.min() + 1, ys.max() - ys.min() + 1)))
378
+ cx0 = max(0, int(xs.min()) - pad)
379
+ cy0 = max(0, int(ys.min()) - pad)
380
+ cx1 = min(width, int(xs.max()) + 1 + pad)
381
+ cy1 = min(height, int(ys.max()) + 1 + pad)
382
+
383
+ crop = image_bgr[cy0:cy1, cx0:cx1]
384
+ crop_mask = mask[cy0:cy1, cx0:cx1]
385
+ crop_h, crop_w = crop.shape[:2]
386
+
387
+ crop_rs = cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA)
388
+ mask_rs = cv2.resize(crop_mask, (size, size), interpolation=cv2.INTER_NEAREST)
389
+
390
+ image_input = cv2.cvtColor(crop_rs, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
391
+ image_input = np.transpose(image_input, (2, 0, 1))[None]
392
+ mask_input = (mask_rs > 127).astype(np.float32)[None, None]
393
+
394
+ output = session.run(None, {image_name: image_input, mask_name: mask_input})[0]
395
+ output = np.asarray(output)[0]
396
+ output = np.transpose(output, (1, 2, 0))
397
+ if float(output.max()) <= 1.5:
398
+ output = output * 255.0
399
+ output = np.clip(output, 0, 255).astype(np.uint8)
400
+ output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
401
+ output_crop = cv2.resize(output_bgr, (crop_w, crop_h), interpolation=cv2.INTER_LINEAR)
402
+
403
+ result = image_bgr.copy()
404
+ region = result[cy0:cy1, cx0:cx1]
405
+ paste = crop_mask > 127
406
+ region[paste] = output_crop[paste]
407
+ result[cy0:cy1, cx0:cx1] = region
408
+ return result
409
+
410
+
411
+ def _migan_forward(model: object, image_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
412
+ import torch
413
+
414
+ image = _norm_img(image_rgb) * 2 - 1
415
+ mask_binary = (mask > 120).astype(np.uint8) * 255
416
+ mask_input = _norm_img(mask_binary)
417
+ image_tensor = torch.from_numpy(image).unsqueeze(0)
418
+ mask_tensor = torch.from_numpy(mask_input).unsqueeze(0)
419
+ erased = image_tensor * (1 - mask_tensor)
420
+ model_input = torch.cat([0.5 - mask_tensor, erased], dim=1)
421
+ with torch.inference_mode():
422
+ output = model(model_input)
423
+ output = (
424
+ (output.permute(0, 2, 3, 1) * 127.5 + 127.5)
425
+ .round()
426
+ .clamp(0, 255)
427
+ .to(torch.uint8)
428
+ )
429
+ result = output[0].cpu().numpy()
430
+ return cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
431
+
432
+
433
+ def _migan_pad_forward(model: object, image_rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
434
+ origin_h, origin_w = image_rgb.shape[:2]
435
+ padded_image = _pad_img_to_modulo(image_rgb, mod=512, square=True, min_size=512)
436
+ padded_mask = _pad_img_to_modulo(mask, mod=512, square=True, min_size=512)
437
+ return _migan_forward(model, padded_image, padded_mask)[:origin_h, :origin_w, :]
438
+
439
+
440
+ def _erase_migan_torchscript(image_bgr: np.ndarray, mask: np.ndarray) -> np.ndarray:
441
+ model = _load_migan_model()
442
+ image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
443
+ if image_rgb.shape[:2] == (512, 512):
444
+ return _migan_pad_forward(model, image_rgb, mask)
445
+
446
+ result = image_rgb[:, :, ::-1].copy()
447
+ for box in _boxes_from_mask(mask):
448
+ crop_img, crop_mask, (left, top, right, bottom) = _crop_box(image_rgb, mask, box, margin=128)
449
+ origin_size = crop_img.shape[:2]
450
+ resized_img = _resize_max_size(crop_img, 512)
451
+ resized_mask = _resize_max_size(crop_mask, 512)
452
+ inpaint = _migan_pad_forward(model, resized_img, resized_mask)
453
+ inpaint = cv2.resize(inpaint, (origin_size[1], origin_size[0]), interpolation=cv2.INTER_CUBIC)
454
+ keep = crop_mask < 127
455
+ inpaint[keep] = crop_img[:, :, ::-1][keep]
456
+ result[top:bottom, left:right, :] = inpaint
457
+ return result
458
+
459
+
460
+ def _normalize_mode(mode: str | None) -> str:
461
+ requested = (mode or DEFAULT_MODE).strip().lower()
462
+ aliases = {
463
+ "lama": "quality",
464
+ "high": "quality",
465
+ "high-quality": "quality",
466
+ "migan": "fast",
467
+ "quick": "fast",
468
+ }
469
+ normalized = aliases.get(requested, requested)
470
+ if normalized not in MODES:
471
+ raise ValueError(f"Unsupported mode: {mode}. Supported modes: {', '.join(MODES)}")
472
+ return normalized
473
+
474
+
475
+ def run_inpaint(image: Image.Image, mask: Image.Image, *, mode: str | None = None) -> tuple[Image.Image, str, float]:
476
+ normalized_mode = _normalize_mode(mode)
477
+ rgb_image = np.array(image.convert("RGB"), copy=True)
478
+ image_bgr = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
479
+ prepared_mask = _prepare_mask(mask, image.size)
480
+
481
+ start = time.time()
482
+ with MODEL_LOCK:
483
+ if normalized_mode == "quality":
484
+ result_bgr = _erase_lama_onnx(image_bgr, prepared_mask)
485
+ elif normalized_mode == "fast":
486
+ result_bgr = _erase_migan_torchscript(image_bgr, prepared_mask)
487
+ else:
488
+ raise ValueError(f"Unsupported mode: {normalized_mode}")
489
+ elapsed = time.time() - start
490
+
491
+ result_rgb = cv2.cvtColor(result_bgr.astype(np.uint8), cv2.COLOR_BGR2RGB)
492
+ return Image.fromarray(result_rgb), normalized_mode, elapsed
493
+
494
+
495
+ def run_remove_doubao(image: Image.Image, *, mode: str | None = None) -> tuple[Image.Image, str, float, dict[str, object]]:
496
+ rgb_image = np.array(image.convert("RGB"), copy=True)
497
+ image_bgr = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
498
+ mask, metadata = generate_doubao_template_bbox_mask(image_bgr)
499
+ mask_image = Image.fromarray(mask, mode="L")
500
+ result, normalized_mode, elapsed = run_inpaint(image, mask_image, mode=mode)
501
+ return result, normalized_mode, elapsed, metadata
502
+
503
+
504
+ @app.on_event("startup")
505
+ def startup() -> None:
506
+ if not LOAD_MODEL_ON_STARTUP:
507
+ return
508
+ if DEFAULT_MODE == "fast":
509
+ _load_migan_model()
510
+ else:
511
+ _load_lama_session()
512
+
513
+
514
+ @app.get("/")
515
+ def root() -> dict[str, object]:
516
+ return health()
517
+
518
+
519
+ @app.get("/health")
520
+ def health() -> dict[str, object]:
521
+ return {
522
+ "status": "ok",
523
+ "mode": "dual-selected-backends",
524
+ "default_mode": DEFAULT_MODE,
525
+ "supported_modes": MODES,
526
+ "loaded": {
527
+ "quality": LAMA_SESSION is not None,
528
+ "fast": MIGAN_MODEL is not None,
529
+ },
530
+ }
531
+
532
+
533
+ @app.post("/inpaint")
534
+ async def inpaint(
535
+ image: UploadFile = File(...),
536
+ mask: UploadFile = File(...),
537
+ mode: str = Form(DEFAULT_MODE),
538
+ ) -> Response:
539
+ try:
540
+ image_bytes = await image.read()
541
+ mask_bytes = await mask.read()
542
+ pil_image = Image.open(io.BytesIO(image_bytes))
543
+ pil_mask = Image.open(io.BytesIO(mask_bytes))
544
+ result, normalized_mode, elapsed = run_inpaint(pil_image, pil_mask, mode=mode)
545
+ output = io.BytesIO()
546
+ result.save(output, format="PNG")
547
+ return Response(
548
+ content=output.getvalue(),
549
+ media_type="image/png",
550
+ headers={
551
+ "X-Inpaint-Mode": normalized_mode,
552
+ "X-Inpaint-Backend": MODES[normalized_mode],
553
+ "X-Inpaint-Elapsed": f"{elapsed:.3f}",
554
+ },
555
+ )
556
+ except Exception as exc:
557
+ raise HTTPException(status_code=500, detail=str(exc)) from exc
558
+
559
+
560
+ @app.post("/remove-doubao")
561
+ async def remove_doubao(
562
+ image: UploadFile = File(...),
563
+ mode: str = Form(DEFAULT_MODE),
564
+ ) -> Response:
565
+ try:
566
+ image_bytes = await image.read()
567
+ pil_image = Image.open(io.BytesIO(image_bytes))
568
+ result, normalized_mode, elapsed, metadata = run_remove_doubao(pil_image, mode=mode)
569
+ output = io.BytesIO()
570
+ result.save(output, format="PNG")
571
+ return Response(
572
+ content=output.getvalue(),
573
+ media_type="image/png",
574
+ headers={
575
+ "X-Inpaint-Mode": normalized_mode,
576
+ "X-Inpaint-Backend": MODES[normalized_mode],
577
+ "X-Inpaint-Elapsed": f"{elapsed:.3f}",
578
+ "X-Doubao-Confidence": str(metadata["confidence"]),
579
+ "X-Doubao-Coverage": str(metadata["coverage"]),
580
+ "X-Doubao-BBox": ",".join(str(v) for v in metadata["bbox"]),
581
+ },
582
+ )
583
+ except Exception as exc:
584
+ raise HTTPException(status_code=500, detail=str(exc)) from exc