update
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""HFUT Model Server 应用包(FastAPI 路由与服务)。"""
|
||||
@@ -0,0 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
# python_server/ 根目录(app/ 的上一级)
|
||||
APP_ROOT = Path(__file__).resolve().parent.parent
|
||||
OUTPUT_DIR = APP_ROOT / "outputs"
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
@@ -0,0 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from .routes import animate, depth, inpaint, meta, segment
|
||||
|
||||
app = FastAPI(title="HFUT Model Server", version="0.1.0")
|
||||
|
||||
app.include_router(meta.router)
|
||||
app.include_router(depth.router)
|
||||
app.include_router(segment.router)
|
||||
app.include_router(inpaint.router)
|
||||
app.include_router(animate.router)
|
||||
@@ -0,0 +1 @@
|
||||
"""HTTP 路由(按领域拆分)。"""
|
||||
@@ -0,0 +1,135 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import shutil
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastapi import APIRouter
|
||||
from PIL import Image
|
||||
|
||||
from ..deps import OUTPUT_DIR
|
||||
from ..schemas import AnimateRequest, CharacterAnimateRequest
|
||||
from ..services.animation_frames import (
|
||||
align_size_to_input,
|
||||
composite_rgba_on_background,
|
||||
list_png_frames,
|
||||
parse_background_color,
|
||||
remove_background_by_color,
|
||||
)
|
||||
from ..services.image_io import b64_to_pil_image, pil_image_to_png_b64
|
||||
from ..services.predictors import get_animation_predictor
|
||||
|
||||
router = APIRouter(tags=["animation"])
|
||||
|
||||
|
||||
@router.post("/animate")
|
||||
def animate(req: AnimateRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "animatediff"
|
||||
predictor = get_animation_predictor(model_name)
|
||||
|
||||
result_path = predictor(
|
||||
prompt=req.prompt,
|
||||
negative_prompt=req.negative_prompt or "",
|
||||
num_inference_steps=req.num_inference_steps,
|
||||
guidance_scale=req.guidance_scale,
|
||||
width=req.width,
|
||||
height=req.height,
|
||||
video_length=req.video_length,
|
||||
seed=req.seed,
|
||||
)
|
||||
|
||||
out_dir = OUTPUT_DIR / "animation"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
out_path = out_dir / f"{model_name}_{ts}.gif"
|
||||
shutil.copy2(result_path, out_path)
|
||||
|
||||
return {"success": True, "output_path": str(out_path)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@router.post("/animate/character_sequence")
|
||||
def animate_character_sequence(req: CharacterAnimateRequest) -> Dict[str, Any]:
|
||||
"""
|
||||
输入透明背景角色 PNG,生成 30fps 角色动画 PNG 序列。
|
||||
|
||||
流程:
|
||||
- 将角色合成到指定纯色背景,作为 AnimateDiff 控制图
|
||||
- 以与输入相近、且对齐到 8 的分辨率生成 PNG 序列
|
||||
- 对每帧按背景色做色键剔除,返回透明背景 PNG 序列给前端
|
||||
"""
|
||||
try:
|
||||
model_name = req.model_name or "animatediff"
|
||||
background_rgb = parse_background_color(req.background_color)
|
||||
|
||||
character = b64_to_pil_image(req.image_b64).convert("RGBA")
|
||||
orig_w, orig_h = character.size
|
||||
run_w, run_h = align_size_to_input(orig_w, orig_h, req.max_side)
|
||||
|
||||
predictor = get_animation_predictor(model_name)
|
||||
|
||||
out_dir = OUTPUT_DIR / "animation" / "character_sequence"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
seq_dir = out_dir / f"{model_name}_{ts}"
|
||||
seq_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
control_path = seq_dir / "control_image.png"
|
||||
control_image = composite_rgba_on_background(
|
||||
character,
|
||||
background_rgb,
|
||||
size=(run_w, run_h),
|
||||
)
|
||||
control_image.save(control_path)
|
||||
|
||||
frame_dir = predictor(
|
||||
prompt=req.prompt,
|
||||
negative_prompt=req.negative_prompt or "",
|
||||
num_inference_steps=req.num_inference_steps,
|
||||
guidance_scale=req.guidance_scale,
|
||||
width=run_w,
|
||||
height=run_h,
|
||||
video_length=req.video_length,
|
||||
seed=req.seed,
|
||||
control_image_path=str(control_path),
|
||||
output_format="png_sequence",
|
||||
)
|
||||
|
||||
frame_paths = list_png_frames(frame_dir)
|
||||
if not frame_paths:
|
||||
raise FileNotFoundError(f"未找到 AnimateDiff 输出 PNG 序列: {frame_dir}")
|
||||
|
||||
frames: list[dict[str, Any]] = []
|
||||
processed_dir = seq_dir / "frames"
|
||||
processed_dir.mkdir(parents=True, exist_ok=True)
|
||||
for idx, frame_path in enumerate(frame_paths):
|
||||
frame = remove_background_by_color(
|
||||
Image.open(frame_path),
|
||||
background_rgb,
|
||||
tolerance=req.background_tolerance,
|
||||
output_size=(orig_w, orig_h),
|
||||
)
|
||||
out_path = processed_dir / f"{idx:04d}.png"
|
||||
frame.save(out_path)
|
||||
frames.append(
|
||||
{
|
||||
"index": idx,
|
||||
"image_b64": pil_image_to_png_b64(frame),
|
||||
"path": str(out_path),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"fps": 30,
|
||||
"frame_count": len(frames),
|
||||
"width": orig_w,
|
||||
"height": orig_h,
|
||||
"background_color": req.background_color,
|
||||
"output_dir": str(processed_dir),
|
||||
"frames": frames,
|
||||
}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e), "frames": []}
|
||||
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import Response
|
||||
|
||||
from ..schemas import DepthRequest
|
||||
from ..services.depth_io import depth_to_png16_bytes
|
||||
from ..services.image_io import b64_to_pil_image
|
||||
from ..services.predictors import get_depth_predictor
|
||||
|
||||
router = APIRouter(tags=["depth"])
|
||||
|
||||
|
||||
@router.post("/depth")
|
||||
def depth(req: DepthRequest):
|
||||
"""
|
||||
计算深度并直接返回二进制 PNG(16-bit 灰度)。
|
||||
|
||||
约束:
|
||||
- 前端不传/不选模型;模型选择写死在后端 config.py
|
||||
- 成功:HTTP 200 + Content-Type: image/png
|
||||
- 失败:HTTP 500,detail 为错误信息
|
||||
"""
|
||||
try:
|
||||
pil = b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
predictor = get_depth_predictor()
|
||||
depth_arr = predictor(pil) # type: ignore[misc]
|
||||
png_bytes = depth_to_png16_bytes(np.asarray(depth_arr))
|
||||
return Response(content=png_bytes, media_type="image/png")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||||
@@ -0,0 +1,47 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from ..deps import OUTPUT_DIR
|
||||
from ..schemas import InpaintRequest
|
||||
from ..services.image_io import b64_to_pil_image, default_half_mask, pil_image_to_png_b64
|
||||
from ..services.predictors import get_inpaint_predictor
|
||||
|
||||
router = APIRouter(tags=["inpaint"])
|
||||
|
||||
|
||||
@router.post("/inpaint")
|
||||
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "sdxl_inpaint"
|
||||
pil = b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
|
||||
if req.mask_b64:
|
||||
mask = b64_to_pil_image(req.mask_b64).convert("L")
|
||||
else:
|
||||
mask = default_half_mask(pil)
|
||||
|
||||
predictor = get_inpaint_predictor(model_name)
|
||||
out = predictor(
|
||||
pil,
|
||||
mask,
|
||||
req.prompt or "",
|
||||
req.negative_prompt or "",
|
||||
strength=req.strength,
|
||||
max_side=req.max_side,
|
||||
)
|
||||
|
||||
out_dir = OUTPUT_DIR / "inpaint"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = out_dir / f"{model_name}_inpaint.png"
|
||||
out.save(out_path)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"output_path": str(out_path),
|
||||
"output_image_b64": pil_image_to_png_b64(out),
|
||||
}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
router = APIRouter(tags=["meta"])
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
def get_models() -> Dict[str, Any]:
|
||||
"""
|
||||
返回一个兼容 Qt 前端的 schema:
|
||||
{
|
||||
"models": {
|
||||
"depth": { "key": { "name": "..."} ... },
|
||||
"segment": { ... },
|
||||
"inpaint": { "key": { "name": "...", "params": [...] } ... }
|
||||
}
|
||||
}
|
||||
"""
|
||||
return {
|
||||
"models": {
|
||||
"depth": {
|
||||
"midas": {"name": "MiDaS (default)"},
|
||||
"zoedepth_n": {"name": "ZoeDepth (ZoeD_N)"},
|
||||
"zoedepth_k": {"name": "ZoeDepth (ZoeD_K)"},
|
||||
"zoedepth_nk": {"name": "ZoeDepth (ZoeD_NK)"},
|
||||
"depth_anything_v2_vits": {"name": "Depth Anything V2 (vits)"},
|
||||
"depth_anything_v2_vitb": {"name": "Depth Anything V2 (vitb)"},
|
||||
"depth_anything_v2_vitl": {"name": "Depth Anything V2 (vitl)"},
|
||||
"depth_anything_v2_vitg": {"name": "Depth Anything V2 (vitg)"},
|
||||
"dpt_large": {"name": "DPT (large)"},
|
||||
"dpt_hybrid": {"name": "DPT (hybrid)"},
|
||||
"midas_dpt_beit_large_512": {"name": "MiDaS (dpt_beit_large_512)"},
|
||||
"midas_dpt_swin2_large_384": {"name": "MiDaS (dpt_swin2_large_384)"},
|
||||
"midas_dpt_swin2_tiny_256": {"name": "MiDaS (dpt_swin2_tiny_256)"},
|
||||
"midas_dpt_levit_224": {"name": "MiDaS (dpt_levit_224)"},
|
||||
},
|
||||
"segment": {
|
||||
"sam": {"name": "SAM (vit_h)"},
|
||||
"mask2former_debug": {"name": "SAM (compat mask2former_debug)"},
|
||||
"mask2former": {"name": "Mask2Former (not implemented)"},
|
||||
},
|
||||
"inpaint": {
|
||||
"copy": {"name": "Copy (no-op)", "params": []},
|
||||
"sdxl_inpaint": {
|
||||
"name": "SDXL Inpaint",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": True},
|
||||
],
|
||||
},
|
||||
"lama": {"name": "LaMa (Erase / Remove Object)", "params": []},
|
||||
"controlnet": {
|
||||
"name": "ControlNet Inpaint (canny)",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": True},
|
||||
],
|
||||
},
|
||||
},
|
||||
"animation": {
|
||||
"animatediff": {
|
||||
"name": "AnimateDiff (Text-to-Video)",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": False},
|
||||
{"id": "negative_prompt", "label": "负向提示词", "optional": True},
|
||||
{"id": "num_inference_steps", "label": "采样步数", "optional": True},
|
||||
{"id": "guidance_scale", "label": "CFG Scale", "optional": True},
|
||||
{"id": "width", "label": "宽度", "optional": True},
|
||||
{"id": "height", "label": "高度", "optional": True},
|
||||
{"id": "video_length", "label": "帧数", "optional": True},
|
||||
{"id": "seed", "label": "随机种子", "optional": True},
|
||||
],
|
||||
},
|
||||
"character_sequence": {
|
||||
"name": "Character PNG Sequence (30 FPS)",
|
||||
"endpoint": "/animate/character_sequence",
|
||||
"params": [
|
||||
{"id": "image_b64", "label": "透明背景角色 PNG", "optional": False},
|
||||
{"id": "prompt", "label": "动画提示词", "optional": False},
|
||||
{"id": "background_color", "label": "纯色背景", "optional": True},
|
||||
{"id": "video_length", "label": "帧数", "optional": True},
|
||||
{"id": "background_tolerance", "label": "背景剔除阈值", "optional": True},
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@router.get("/health")
|
||||
def health() -> Dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
@@ -0,0 +1,84 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
import numpy as np
|
||||
from fastapi import APIRouter
|
||||
from PIL import Image
|
||||
|
||||
from ..deps import OUTPUT_DIR
|
||||
from ..schemas import SamPromptSegmentRequest, SegmentRequest
|
||||
from ..services.image_io import b64_to_pil_image
|
||||
from ..services.predictors import get_seg_predictor
|
||||
from model.Seg.seg_loader import expand_mask, mask_to_contour_xy, run_sam_prompt
|
||||
|
||||
router = APIRouter(tags=["segment"])
|
||||
|
||||
|
||||
@router.post("/segment")
|
||||
def segment(req: SegmentRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "sam"
|
||||
pil = b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
rgb = np.array(pil, dtype=np.uint8)
|
||||
|
||||
predictor = get_seg_predictor(model_name)
|
||||
label_map = predictor(rgb).astype(np.int32)
|
||||
|
||||
out_dir = OUTPUT_DIR / "segment"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = out_dir / f"{model_name}_label.png"
|
||||
Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)
|
||||
|
||||
return {"success": True, "label_path": str(out_path)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@router.post("/segment/sam_prompt")
|
||||
def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
|
||||
"""
|
||||
交互式 SAM:裁剪图 + 点/框提示,返回掩膜外轮廓点列(裁剪像素坐标)。
|
||||
"""
|
||||
try:
|
||||
pil = b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
rgb = np.array(pil, dtype=np.uint8)
|
||||
h, w = rgb.shape[0], rgb.shape[1]
|
||||
|
||||
if req.overlay_b64:
|
||||
ov = b64_to_pil_image(req.overlay_b64)
|
||||
if ov.size != (w, h):
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
|
||||
"contour": [],
|
||||
}
|
||||
|
||||
if len(req.point_coords) != len(req.point_labels):
|
||||
return {
|
||||
"success": False,
|
||||
"error": "point_coords 与 point_labels 长度不一致",
|
||||
"contour": [],
|
||||
}
|
||||
if len(req.point_coords) < 1:
|
||||
return {"success": False, "error": "至少需要一个提示点", "contour": []}
|
||||
|
||||
pc = np.array(req.point_coords, dtype=np.float32)
|
||||
if pc.ndim != 2 or pc.shape[1] != 2:
|
||||
return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}
|
||||
|
||||
pl = np.array(req.point_labels, dtype=np.int64)
|
||||
box = np.array(req.box_xyxy, dtype=np.float32)
|
||||
|
||||
mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
|
||||
if not np.any(mask):
|
||||
return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}
|
||||
|
||||
mask = expand_mask(mask, int(req.expand_px))
|
||||
contour = mask_to_contour_xy(mask, epsilon_px=2.0)
|
||||
if len(contour) < 3:
|
||||
return {"success": False, "error": "轮廓点数不足", "contour": []}
|
||||
|
||||
return {"success": True, "contour": contour, "error": None}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e), "contour": []}
|
||||
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ImageInput(BaseModel):
|
||||
image_b64: str = Field(..., description="PNG/JPG 编码后的 base64(不含 data: 前缀)")
|
||||
model_name: Optional[str] = Field(None, description="模型 key(来自 /models)")
|
||||
|
||||
|
||||
class DepthRequest(ImageInput):
|
||||
pass
|
||||
|
||||
|
||||
class SegmentRequest(ImageInput):
|
||||
pass
|
||||
|
||||
|
||||
class SamPromptSegmentRequest(BaseModel):
|
||||
image_b64: str = Field(..., description="裁剪后的 RGB 图 base64(PNG/JPG)")
|
||||
overlay_b64: Optional[str] = Field(
|
||||
None,
|
||||
description="与裁剪同尺寸的标记叠加 PNG base64(可选;当前用于校验尺寸一致)",
|
||||
)
|
||||
point_coords: list[list[float]] = Field(
|
||||
...,
|
||||
description="裁剪坐标系下的提示点 [[x,y], ...]",
|
||||
)
|
||||
point_labels: list[int] = Field(
|
||||
...,
|
||||
description="与 point_coords 等长:1=前景,0=背景",
|
||||
)
|
||||
box_xyxy: list[float] = Field(
|
||||
...,
|
||||
description="裁剪内笔画紧包围盒 [x1,y1,x2,y2](像素)",
|
||||
min_length=4,
|
||||
max_length=4,
|
||||
)
|
||||
expand_px: int = Field(
|
||||
2,
|
||||
ge=0,
|
||||
le=20,
|
||||
description="对输出 mask 做轻微膨胀的像素半径(0 表示不扩大;建议 1~3)",
|
||||
)
|
||||
|
||||
|
||||
class InpaintRequest(ImageInput):
|
||||
prompt: Optional[str] = Field("", description="补全 prompt")
|
||||
strength: float = Field(0.8, ge=0.0, le=1.0)
|
||||
negative_prompt: Optional[str] = Field("", description="负向 prompt")
|
||||
mask_b64: Optional[str] = Field(None, description="mask PNG base64(可选)")
|
||||
max_side: int = Field(1024, ge=128, le=2048)
|
||||
|
||||
|
||||
class AnimateRequest(BaseModel):
|
||||
model_name: Optional[str] = Field(None, description="模型 key(来自 /models)")
|
||||
prompt: str = Field(..., description="文本提示词")
|
||||
negative_prompt: Optional[str] = Field("", description="负向提示词")
|
||||
num_inference_steps: int = Field(25, ge=1, le=200)
|
||||
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
|
||||
width: int = Field(512, ge=128, le=2048)
|
||||
height: int = Field(512, ge=128, le=2048)
|
||||
video_length: int = Field(16, ge=1, le=128)
|
||||
seed: int = Field(-1, description="-1 表示随机种子")
|
||||
|
||||
|
||||
class CharacterAnimateRequest(BaseModel):
|
||||
image_b64: str = Field(..., description="透明背景角色 PNG 的 base64(不含 data: 前缀)")
|
||||
model_name: Optional[str] = Field(None, description="模型 key(来自 /models)")
|
||||
prompt: str = Field(..., description="角色动画提示词")
|
||||
negative_prompt: Optional[str] = Field("", description="负向提示词")
|
||||
background_color: str = Field(
|
||||
"#00FF00",
|
||||
description="生成时使用的纯色背景,支持 #RRGGBB 或 R,G,B",
|
||||
)
|
||||
background_tolerance: int = Field(
|
||||
40,
|
||||
ge=0,
|
||||
le=255,
|
||||
description="自动剔除背景的 RGB 色差阈值,越大剔除越激进",
|
||||
)
|
||||
num_inference_steps: int = Field(25, ge=1, le=200)
|
||||
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
|
||||
video_length: int = Field(16, ge=1, le=128, description="前端请求生成的帧数")
|
||||
seed: int = Field(-1, description="-1 表示随机种子")
|
||||
max_side: int = Field(
|
||||
768,
|
||||
ge=128,
|
||||
le=2048,
|
||||
description="推理分辨率长边上限;会保持输入比例并对齐到 8 的倍数",
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
"""业务服务层(预测器缓存、图像编解码等)。"""
|
||||
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def parse_background_color(value: str) -> tuple[int, int, int]:
|
||||
text = value.strip()
|
||||
if text.startswith("#"):
|
||||
hex_text = text[1:]
|
||||
if len(hex_text) != 6:
|
||||
raise ValueError("background_color 使用 #RRGGBB 时必须是 6 位十六进制")
|
||||
try:
|
||||
return tuple(int(hex_text[i : i + 2], 16) for i in (0, 2, 4)) # type: ignore[return-value]
|
||||
except ValueError as e:
|
||||
raise ValueError(f"background_color 不是有效十六进制颜色: {value}") from e
|
||||
|
||||
parts = [p.strip() for p in text.split(",")]
|
||||
if len(parts) != 3:
|
||||
raise ValueError("background_color 必须为 #RRGGBB 或 R,G,B")
|
||||
try:
|
||||
rgb = tuple(int(p) for p in parts)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"background_color 不是有效 RGB 颜色: {value}") from e
|
||||
if any(c < 0 or c > 255 for c in rgb):
|
||||
raise ValueError("background_color 的 RGB 分量必须在 0~255")
|
||||
return rgb # type: ignore[return-value]
|
||||
|
||||
|
||||
def align_size_to_input(width: int, height: int, max_side: int) -> tuple[int, int]:
|
||||
run_w, run_h = width, height
|
||||
if max(width, height) > max_side:
|
||||
scale = max_side / float(max(width, height))
|
||||
run_w = int(round(width * scale))
|
||||
run_h = int(round(height * scale))
|
||||
|
||||
run_w = max(8, run_w - (run_w % 8))
|
||||
run_h = max(8, run_h - (run_h % 8))
|
||||
return run_w, run_h
|
||||
|
||||
|
||||
def composite_rgba_on_background(
|
||||
image: Image.Image,
|
||||
background_rgb: tuple[int, int, int],
|
||||
size: tuple[int, int],
|
||||
) -> Image.Image:
|
||||
rgba = image.convert("RGBA")
|
||||
if rgba.size != size:
|
||||
rgba = rgba.resize(size, resample=Image.BICUBIC)
|
||||
|
||||
bg = Image.new("RGB", size, background_rgb)
|
||||
alpha = rgba.getchannel("A")
|
||||
bg.paste(rgba.convert("RGB"), mask=alpha)
|
||||
return bg
|
||||
|
||||
|
||||
def remove_background_by_color(
|
||||
image: Image.Image,
|
||||
background_rgb: tuple[int, int, int],
|
||||
tolerance: int,
|
||||
output_size: tuple[int, int] | None = None,
|
||||
) -> Image.Image:
|
||||
rgba = image.convert("RGBA")
|
||||
if output_size is not None and rgba.size != output_size:
|
||||
rgba = rgba.resize(output_size, resample=Image.BICUBIC)
|
||||
|
||||
arr = np.asarray(rgba, dtype=np.uint8).copy()
|
||||
rgb = arr[:, :, :3].astype(np.int16)
|
||||
bg = np.array(background_rgb, dtype=np.int16).reshape(1, 1, 3)
|
||||
diff = np.abs(rgb - bg)
|
||||
bg_like = np.max(diff, axis=2) <= int(tolerance)
|
||||
bg_mask = _edge_connected_mask(bg_like)
|
||||
arr[bg_mask, 3] = 0
|
||||
return Image.fromarray(arr, mode="RGBA")
|
||||
|
||||
|
||||
def list_png_frames(frame_dir: Path) -> list[Path]:
|
||||
frames = [p for p in frame_dir.iterdir() if p.is_file() and p.suffix.lower() == ".png"]
|
||||
return sorted(frames, key=lambda p: p.name)
|
||||
|
||||
|
||||
def _edge_connected_mask(mask: np.ndarray) -> np.ndarray:
|
||||
h, w = mask.shape
|
||||
visited = np.zeros((h, w), dtype=bool)
|
||||
queue: deque[tuple[int, int]] = deque()
|
||||
|
||||
for x in range(w):
|
||||
if mask[0, x]:
|
||||
queue.append((0, x))
|
||||
if mask[h - 1, x]:
|
||||
queue.append((h - 1, x))
|
||||
for y in range(h):
|
||||
if mask[y, 0]:
|
||||
queue.append((y, 0))
|
||||
if mask[y, w - 1]:
|
||||
queue.append((y, w - 1))
|
||||
|
||||
while queue:
|
||||
y, x = queue.popleft()
|
||||
if visited[y, x] or not mask[y, x]:
|
||||
continue
|
||||
visited[y, x] = True
|
||||
if y > 0:
|
||||
queue.append((y - 1, x))
|
||||
if y + 1 < h:
|
||||
queue.append((y + 1, x))
|
||||
if x > 0:
|
||||
queue.append((y, x - 1))
|
||||
if x + 1 < w:
|
||||
queue.append((y, x + 1))
|
||||
|
||||
return visited
|
||||
@@ -0,0 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def depth_to_png16_bytes(depth: np.ndarray) -> bytes:
|
||||
depth = np.asarray(depth, dtype=np.float32)
|
||||
dmin = float(depth.min())
|
||||
dmax = float(depth.max())
|
||||
if dmax > dmin:
|
||||
norm = (depth - dmin) / (dmax - dmin)
|
||||
else:
|
||||
norm = np.zeros_like(depth, dtype=np.float32)
|
||||
# 前后端约定:最远=255,最近=0(8-bit)
|
||||
u8 = ((1.0 - norm) * 255.0).clip(0, 255).astype(np.uint8)
|
||||
img = Image.fromarray(u8, mode="L")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
@@ -0,0 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def b64_to_pil_image(b64: str) -> Image.Image:
|
||||
raw = base64.b64decode(b64)
|
||||
return Image.open(io.BytesIO(raw))
|
||||
|
||||
|
||||
def pil_image_to_png_b64(img: Image.Image) -> str:
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def default_half_mask(img: Image.Image) -> Image.Image:
|
||||
w, h = img.size
|
||||
mask = Image.new("L", (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([w // 2, 0, w, h], fill=255)
|
||||
return mask
|
||||
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from config_loader import load_app_config, get_depth_backend_from_app
|
||||
from model.Animation.animation_loader import (
|
||||
UnifiedAnimationConfig,
|
||||
AnimationBackend,
|
||||
build_animation_predictor,
|
||||
)
|
||||
from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
|
||||
from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
|
||||
from model.Seg.seg_loader import UnifiedSegConfig, SegBackend, build_seg_predictor
|
||||
|
||||
_depth_predictor = None
|
||||
_depth_backend: DepthBackend | None = None
|
||||
|
||||
_seg_cache: dict[str, Any] = {}
|
||||
|
||||
_inpaint_cache: dict[str, Any] = {}
|
||||
|
||||
_animation_cache: dict[str, Any] = {}
|
||||
|
||||
|
||||
def ensure_depth_predictor() -> None:
|
||||
global _depth_predictor, _depth_backend
|
||||
if _depth_predictor is not None and _depth_backend is not None:
|
||||
return
|
||||
|
||||
app_cfg = load_app_config()
|
||||
backend_str = get_depth_backend_from_app(app_cfg)
|
||||
try:
|
||||
backend = DepthBackend(backend_str)
|
||||
except Exception as e:
|
||||
raise ValueError(f"config.py 中 depth.backend 不合法: {backend_str}") from e
|
||||
|
||||
_depth_predictor, _depth_backend = build_depth_predictor(UnifiedDepthConfig(backend=backend))
|
||||
|
||||
|
||||
def get_depth_predictor():
|
||||
ensure_depth_predictor()
|
||||
return _depth_predictor
|
||||
|
||||
|
||||
def get_seg_predictor(model_name: str):
|
||||
if model_name == "mask2former_debug":
|
||||
model_name = "sam"
|
||||
|
||||
if model_name in _seg_cache:
|
||||
return _seg_cache[model_name]
|
||||
|
||||
if model_name == "sam":
|
||||
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.SAM))
|
||||
_seg_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
if model_name == "mask2former":
|
||||
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.MASK2FORMER))
|
||||
_seg_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 segment model_name: {model_name}")
|
||||
|
||||
|
||||
def get_inpaint_predictor(model_name: str):
|
||||
if model_name in _inpaint_cache:
|
||||
return _inpaint_cache[model_name]
|
||||
|
||||
if model_name == "copy":
|
||||
|
||||
def _copy(image: Image.Image, *_args, **_kwargs) -> Image.Image:
|
||||
return image.convert("RGB")
|
||||
|
||||
_inpaint_cache[model_name] = _copy
|
||||
return _copy
|
||||
|
||||
if model_name == "sdxl_inpaint":
|
||||
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.SDXL_INPAINT))
|
||||
_inpaint_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
if model_name == "controlnet":
|
||||
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.CONTROLNET))
|
||||
_inpaint_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
if model_name == "lama":
|
||||
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.LAMA))
|
||||
_inpaint_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 inpaint model_name: {model_name}")
|
||||
|
||||
|
||||
def get_animation_predictor(model_name: str):
|
||||
if model_name in _animation_cache:
|
||||
return _animation_cache[model_name]
|
||||
|
||||
if model_name == "animatediff":
|
||||
pred, _ = build_animation_predictor(
|
||||
UnifiedAnimationConfig(backend=AnimationBackend.ANIMATEDIFF)
|
||||
)
|
||||
_animation_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 animation model_name: {model_name}")
|
||||
@@ -5,7 +5,8 @@ from __future__ import annotations
|
||||
|
||||
当前支持:
|
||||
- SDXL Inpaint(diffusers AutoPipelineForInpainting)
|
||||
- ControlNet(占位,暂未统一封装)
|
||||
- ControlNet(Stable Diffusion + ControlNet Inpaint)
|
||||
- LaMa(torchscript big-lama,本地封装)
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
@@ -26,6 +27,7 @@ from config_loader import (
|
||||
class InpaintBackend(str, Enum):
|
||||
SDXL_INPAINT = "sdxl_inpaint"
|
||||
CONTROLNET = "controlnet"
|
||||
LAMA = "lama"
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -278,11 +280,13 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
|
||||
image_run = image
|
||||
mask_run = mask
|
||||
|
||||
# control image:使用 canny 边缘作为约束(最通用)
|
||||
# control image:使用 Canny 边缘作为条件(更通用、更稳定)。
|
||||
# 备注:当前项目的 ControlNet 权重未必是 inpaint 专用模型,
|
||||
# 若直接把洞内置零作为条件,容易把输出“拉黑”。先回退到稳定方案。
|
||||
rgb = np.array(image_run, dtype=np.uint8)
|
||||
edges = cv2.Canny(rgb, 100, 200)
|
||||
edges3 = np.stack([edges, edges, edges], axis=-1)
|
||||
control_image = Image.fromarray(edges3)
|
||||
control_image = Image.fromarray(edges3, mode="RGB")
|
||||
|
||||
if device == "cuda":
|
||||
try:
|
||||
@@ -312,6 +316,37 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
|
||||
return _predict
|
||||
|
||||
|
||||
def _make_lama_predictor(cfg: UnifiedInpaintConfig):
|
||||
import torch
|
||||
|
||||
app_cfg = load_app_config()
|
||||
device = cfg.device
|
||||
if device is None:
|
||||
device = app_cfg.inpaint.device
|
||||
device = "cuda" if device.startswith("cuda") and torch.cuda.is_available() else "cpu"
|
||||
|
||||
from .lama_torchscript import LaMaTorchscript
|
||||
|
||||
lama = LaMaTorchscript(device=device)
|
||||
|
||||
def _predict(
|
||||
image: Image.Image,
|
||||
mask: Image.Image,
|
||||
_prompt: str,
|
||||
_negative_prompt: str = "",
|
||||
strength: float = 0.8,
|
||||
guidance_scale: float = 7.5,
|
||||
num_inference_steps: int = 30,
|
||||
max_side: int = 1024,
|
||||
**_kwargs,
|
||||
) -> Image.Image:
|
||||
# LaMa 不使用 prompt/negative_prompt,也不使用 strength/steps 等扩散参数
|
||||
_ = (strength, guidance_scale, num_inference_steps, max_side)
|
||||
return lama(image, mask)
|
||||
|
||||
return _predict
|
||||
|
||||
|
||||
def build_inpaint_predictor(
|
||||
cfg: UnifiedInpaintConfig | None = None,
|
||||
) -> tuple[Callable[..., Image.Image], InpaintBackend]:
|
||||
@@ -326,6 +361,9 @@ def build_inpaint_predictor(
|
||||
if cfg.backend == InpaintBackend.CONTROLNET:
|
||||
return _make_controlnet_predictor(cfg), InpaintBackend.CONTROLNET
|
||||
|
||||
if cfg.backend == InpaintBackend.LAMA:
|
||||
return _make_lama_predictor(cfg), InpaintBackend.LAMA
|
||||
|
||||
raise ValueError(f"不支持的补全后端: {cfg.backend}")
|
||||
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -95,6 +95,43 @@ def run_sam_prompt(
|
||||
return m
|
||||
|
||||
|
||||
def expand_mask(mask_bool: np.ndarray, expand_px: int = 0) -> np.ndarray:
|
||||
"""
|
||||
轻微扩大二值 mask(用于消除边界过软/过细带来的“漏边”观感)。
|
||||
|
||||
- expand_px <= 0: 原样返回
|
||||
- expand_px > 0: 做 dilation(膨胀)expand_px 像素级别
|
||||
"""
|
||||
mask_bool = np.asarray(mask_bool, dtype=bool)
|
||||
if expand_px <= 0:
|
||||
return mask_bool
|
||||
|
||||
# 优先用 OpenCV(通常更快)
|
||||
try:
|
||||
import cv2 # type: ignore[import]
|
||||
|
||||
u8 = (mask_bool.astype(np.uint8) * 255)
|
||||
k = int(expand_px) * 2 + 1
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
out = cv2.dilate(u8, kernel, iterations=1)
|
||||
return out > 127
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 回退到 scipy(若存在)
|
||||
try:
|
||||
from scipy.ndimage import binary_dilation # type: ignore[import]
|
||||
|
||||
# 使用近似圆形结构元素
|
||||
r = int(expand_px)
|
||||
yy, xx = np.ogrid[-r : r + 1, -r : r + 1]
|
||||
selem = (xx * xx + yy * yy) <= (r * r)
|
||||
return binary_dilation(mask_bool, structure=selem)
|
||||
except Exception:
|
||||
# 再回退:无依赖时就不扩大(保持可用性)
|
||||
return mask_bool
|
||||
|
||||
|
||||
def mask_to_contour_xy(
|
||||
mask_bool: np.ndarray,
|
||||
epsilon_px: float = 2.0,
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
accelerate==1.13.0
|
||||
annotated-doc==0.0.4
|
||||
annotated-types==0.7.0
|
||||
antlr4-python3-runtime==4.9.3
|
||||
anyio==4.12.1
|
||||
certifi==2026.2.25
|
||||
charset-normalizer==3.4.5
|
||||
click==8.3.1
|
||||
cuda-bindings==12.9.4
|
||||
cuda-pathfinder==1.4.1
|
||||
diffusers==0.37.0
|
||||
einops==0.8.2
|
||||
fastapi==0.135.1
|
||||
filelock==3.25.0
|
||||
fsspec==2026.2.0
|
||||
h11==0.16.0
|
||||
hf-xet==1.3.2
|
||||
httpcore==1.0.9
|
||||
httptools==0.7.1
|
||||
httpx==0.28.1
|
||||
huggingface_hub==1.6.0
|
||||
idna==3.11
|
||||
ImageIO==2.37.3
|
||||
importlib_metadata==8.7.1
|
||||
Jinja2==3.1.6
|
||||
markdown-it-py==4.0.0
|
||||
MarkupSafe==3.0.3
|
||||
mdurl==0.1.2
|
||||
mpmath==1.3.0
|
||||
networkx==3.6.1
|
||||
numpy==2.4.3
|
||||
nvidia-cublas-cu12==12.8.4.1
|
||||
nvidia-cuda-cupti-cu12==12.8.90
|
||||
nvidia-cuda-nvrtc-cu12==12.8.93
|
||||
nvidia-cuda-runtime-cu12==12.8.90
|
||||
nvidia-cudnn-cu12==9.10.2.21
|
||||
nvidia-cufft-cu12==11.3.3.83
|
||||
nvidia-cufile-cu12==1.13.1.3
|
||||
nvidia-curand-cu12==10.3.9.90
|
||||
nvidia-cusolver-cu12==11.7.3.90
|
||||
nvidia-cusparse-cu12==12.5.8.93
|
||||
nvidia-cusparselt-cu12==0.7.1
|
||||
nvidia-nccl-cu12==2.27.5
|
||||
nvidia-nvjitlink-cu12==12.8.93
|
||||
nvidia-nvshmem-cu12==3.4.5
|
||||
nvidia-nvtx-cu12==12.8.90
|
||||
omegaconf==2.3.0
|
||||
opencv-python==4.13.0.92
|
||||
packaging==26.0
|
||||
pillow==12.1.1
|
||||
psutil==7.2.2
|
||||
pydantic==2.12.5
|
||||
pydantic_core==2.41.5
|
||||
Pygments==2.19.2
|
||||
python-dotenv==1.2.2
|
||||
PyYAML==6.0.3
|
||||
regex==2026.2.28
|
||||
requests==2.32.5
|
||||
rich==14.3.3
|
||||
safetensors==0.7.0
|
||||
scipy==1.17.1
|
||||
setuptools==82.0.0
|
||||
shellingham==1.5.4
|
||||
starlette==0.52.1
|
||||
sympy==1.14.0
|
||||
timm==0.6.13
|
||||
tokenizers==0.22.2
|
||||
torch==2.10.0
|
||||
torchvision==0.25.0
|
||||
tqdm==4.67.3
|
||||
transformers==5.3.0
|
||||
triton==3.6.0
|
||||
typer==0.24.1
|
||||
typing-inspection==0.4.2
|
||||
typing_extensions==4.15.0
|
||||
urllib3==2.6.3
|
||||
uvicorn==0.41.0
|
||||
uvloop==0.22.1
|
||||
watchfiles==1.1.1
|
||||
websockets==16.0
|
||||
xformers==0.0.35
|
||||
zipp==3.23.0
|
||||
+7
-477
@@ -1,488 +1,18 @@
|
||||
"""
|
||||
source .venv/bin/activate # 仓库根 hfut-bishe/.venv
|
||||
cd python_server && uvicorn app.main:app --host 0.0.0.0 --port 8000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import datetime
|
||||
import io
|
||||
import os
|
||||
import shutil
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import numpy as np
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import Response
|
||||
from pydantic import BaseModel, Field
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from config_loader import load_app_config, get_depth_backend_from_app
|
||||
from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
|
||||
|
||||
from model.Seg.seg_loader import (
|
||||
UnifiedSegConfig,
|
||||
SegBackend,
|
||||
build_seg_predictor,
|
||||
mask_to_contour_xy,
|
||||
run_sam_prompt,
|
||||
)
|
||||
from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
|
||||
from model.Animation.animation_loader import (
|
||||
UnifiedAnimationConfig,
|
||||
AnimationBackend,
|
||||
build_animation_predictor,
|
||||
)
|
||||
|
||||
|
||||
APP_ROOT = Path(__file__).resolve().parent
|
||||
OUTPUT_DIR = APP_ROOT / "outputs"
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
app = FastAPI(title="HFUT Model Server", version="0.1.0")
|
||||
|
||||
|
||||
class ImageInput(BaseModel):
|
||||
image_b64: str = Field(..., description="PNG/JPG 编码后的 base64(不含 data: 前缀)")
|
||||
model_name: Optional[str] = Field(None, description="模型 key(来自 /models)")
|
||||
|
||||
|
||||
class DepthRequest(ImageInput):
|
||||
pass
|
||||
|
||||
|
||||
class SegmentRequest(ImageInput):
|
||||
pass
|
||||
|
||||
|
||||
class SamPromptSegmentRequest(BaseModel):
|
||||
image_b64: str = Field(..., description="裁剪后的 RGB 图 base64(PNG/JPG)")
|
||||
overlay_b64: Optional[str] = Field(
|
||||
None,
|
||||
description="与裁剪同尺寸的标记叠加 PNG base64(可选;当前用于校验尺寸一致)",
|
||||
)
|
||||
point_coords: list[list[float]] = Field(
|
||||
...,
|
||||
description="裁剪坐标系下的提示点 [[x,y], ...]",
|
||||
)
|
||||
point_labels: list[int] = Field(
|
||||
...,
|
||||
description="与 point_coords 等长:1=前景,0=背景",
|
||||
)
|
||||
box_xyxy: list[float] = Field(
|
||||
...,
|
||||
description="裁剪内笔画紧包围盒 [x1,y1,x2,y2](像素)",
|
||||
min_length=4,
|
||||
max_length=4,
|
||||
)
|
||||
|
||||
|
||||
class InpaintRequest(ImageInput):
|
||||
prompt: Optional[str] = Field("", description="补全 prompt")
|
||||
strength: float = Field(0.8, ge=0.0, le=1.0)
|
||||
negative_prompt: Optional[str] = Field("", description="负向 prompt")
|
||||
# 可选 mask(白色区域为重绘)
|
||||
mask_b64: Optional[str] = Field(None, description="mask PNG base64(可选)")
|
||||
# 推理缩放上限(避免 OOM)
|
||||
max_side: int = Field(1024, ge=128, le=2048)
|
||||
|
||||
|
||||
class AnimateRequest(BaseModel):
|
||||
model_name: Optional[str] = Field(None, description="模型 key(来自 /models)")
|
||||
prompt: str = Field(..., description="文本提示词")
|
||||
negative_prompt: Optional[str] = Field("", description="负向提示词")
|
||||
num_inference_steps: int = Field(25, ge=1, le=200)
|
||||
guidance_scale: float = Field(8.0, ge=0.0, le=30.0)
|
||||
width: int = Field(512, ge=128, le=2048)
|
||||
height: int = Field(512, ge=128, le=2048)
|
||||
video_length: int = Field(16, ge=1, le=128)
|
||||
seed: int = Field(-1, description="-1 表示随机种子")
|
||||
|
||||
|
||||
def _b64_to_pil_image(b64: str) -> Image.Image:
|
||||
raw = base64.b64decode(b64)
|
||||
return Image.open(io.BytesIO(raw))
|
||||
|
||||
|
||||
def _pil_image_to_png_b64(img: Image.Image) -> str:
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def _depth_to_png16_b64(depth: np.ndarray) -> str:
|
||||
depth = np.asarray(depth, dtype=np.float32)
|
||||
dmin = float(depth.min())
|
||||
dmax = float(depth.max())
|
||||
if dmax > dmin:
|
||||
norm = (depth - dmin) / (dmax - dmin)
|
||||
else:
|
||||
norm = np.zeros_like(depth, dtype=np.float32)
|
||||
u16 = (norm * 65535.0).clip(0, 65535).astype(np.uint16)
|
||||
img = Image.fromarray(u16, mode="I;16")
|
||||
return _pil_image_to_png_b64(img)
|
||||
|
||||
def _depth_to_png16_bytes(depth: np.ndarray) -> bytes:
|
||||
depth = np.asarray(depth, dtype=np.float32)
|
||||
dmin = float(depth.min())
|
||||
dmax = float(depth.max())
|
||||
if dmax > dmin:
|
||||
norm = (depth - dmin) / (dmax - dmin)
|
||||
else:
|
||||
norm = np.zeros_like(depth, dtype=np.float32)
|
||||
# 前后端约定:最远=255,最近=0(8-bit)
|
||||
u8 = ((1.0 - norm) * 255.0).clip(0, 255).astype(np.uint8)
|
||||
img = Image.fromarray(u8, mode="L")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _default_half_mask(img: Image.Image) -> Image.Image:
|
||||
w, h = img.size
|
||||
mask = Image.new("L", (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([w // 2, 0, w, h], fill=255)
|
||||
return mask
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# /models(给前端/GUI 使用)
|
||||
# -----------------------------
|
||||
|
||||
|
||||
@app.get("/models")
|
||||
def get_models() -> Dict[str, Any]:
|
||||
"""
|
||||
返回一个兼容 Qt 前端的 schema:
|
||||
{
|
||||
"models": {
|
||||
"depth": { "key": { "name": "..."} ... },
|
||||
"segment": { ... },
|
||||
"inpaint": { "key": { "name": "...", "params": [...] } ... }
|
||||
}
|
||||
}
|
||||
"""
|
||||
return {
|
||||
"models": {
|
||||
"depth": {
|
||||
# 兼容旧配置默认值
|
||||
"midas": {"name": "MiDaS (default)"},
|
||||
"zoedepth_n": {"name": "ZoeDepth (ZoeD_N)"},
|
||||
"zoedepth_k": {"name": "ZoeDepth (ZoeD_K)"},
|
||||
"zoedepth_nk": {"name": "ZoeDepth (ZoeD_NK)"},
|
||||
"depth_anything_v2_vits": {"name": "Depth Anything V2 (vits)"},
|
||||
"depth_anything_v2_vitb": {"name": "Depth Anything V2 (vitb)"},
|
||||
"depth_anything_v2_vitl": {"name": "Depth Anything V2 (vitl)"},
|
||||
"depth_anything_v2_vitg": {"name": "Depth Anything V2 (vitg)"},
|
||||
"dpt_large": {"name": "DPT (large)"},
|
||||
"dpt_hybrid": {"name": "DPT (hybrid)"},
|
||||
"midas_dpt_beit_large_512": {"name": "MiDaS (dpt_beit_large_512)"},
|
||||
"midas_dpt_swin2_large_384": {"name": "MiDaS (dpt_swin2_large_384)"},
|
||||
"midas_dpt_swin2_tiny_256": {"name": "MiDaS (dpt_swin2_tiny_256)"},
|
||||
"midas_dpt_levit_224": {"name": "MiDaS (dpt_levit_224)"},
|
||||
},
|
||||
"segment": {
|
||||
"sam": {"name": "SAM (vit_h)"},
|
||||
# 兼容旧配置默认值
|
||||
"mask2former_debug": {"name": "SAM (compat mask2former_debug)"},
|
||||
"mask2former": {"name": "Mask2Former (not implemented)"},
|
||||
},
|
||||
"inpaint": {
|
||||
# 兼容旧配置默认值:copy 表示不做补全
|
||||
"copy": {"name": "Copy (no-op)", "params": []},
|
||||
"sdxl_inpaint": {
|
||||
"name": "SDXL Inpaint",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": True},
|
||||
],
|
||||
},
|
||||
"controlnet": {
|
||||
"name": "ControlNet Inpaint (canny)",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": True},
|
||||
],
|
||||
},
|
||||
},
|
||||
"animation": {
|
||||
"animatediff": {
|
||||
"name": "AnimateDiff (Text-to-Video)",
|
||||
"params": [
|
||||
{"id": "prompt", "label": "提示词", "optional": False},
|
||||
{"id": "negative_prompt", "label": "负向提示词", "optional": True},
|
||||
{"id": "num_inference_steps", "label": "采样步数", "optional": True},
|
||||
{"id": "guidance_scale", "label": "CFG Scale", "optional": True},
|
||||
{"id": "width", "label": "宽度", "optional": True},
|
||||
{"id": "height", "label": "高度", "optional": True},
|
||||
{"id": "video_length", "label": "帧数", "optional": True},
|
||||
{"id": "seed", "label": "随机种子", "optional": True},
|
||||
],
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Depth
|
||||
# -----------------------------
|
||||
|
||||
|
||||
_depth_predictor = None
|
||||
_depth_backend: DepthBackend | None = None
|
||||
|
||||
|
||||
def _ensure_depth_predictor() -> None:
|
||||
global _depth_predictor, _depth_backend
|
||||
if _depth_predictor is not None and _depth_backend is not None:
|
||||
return
|
||||
|
||||
app_cfg = load_app_config()
|
||||
backend_str = get_depth_backend_from_app(app_cfg)
|
||||
try:
|
||||
backend = DepthBackend(backend_str)
|
||||
except Exception as e:
|
||||
raise ValueError(f"config.py 中 depth.backend 不合法: {backend_str}") from e
|
||||
|
||||
_depth_predictor, _depth_backend = build_depth_predictor(UnifiedDepthConfig(backend=backend))
|
||||
|
||||
|
||||
@app.post("/depth")
|
||||
def depth(req: DepthRequest):
|
||||
"""
|
||||
计算深度并直接返回二进制 PNG(16-bit 灰度)。
|
||||
|
||||
约束:
|
||||
- 前端不传/不选模型;模型选择写死在后端 config.py
|
||||
- 成功:HTTP 200 + Content-Type: image/png
|
||||
- 失败:HTTP 500,detail 为错误信息
|
||||
"""
|
||||
try:
|
||||
_ensure_depth_predictor()
|
||||
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
depth_arr = _depth_predictor(pil) # type: ignore[misc]
|
||||
png_bytes = _depth_to_png16_bytes(np.asarray(depth_arr))
|
||||
return Response(content=png_bytes, media_type="image/png")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Segment
|
||||
# -----------------------------
|
||||
|
||||
|
||||
_seg_cache: Dict[str, Any] = {}
|
||||
|
||||
|
||||
def _get_seg_predictor(model_name: str):
|
||||
if model_name in _seg_cache:
|
||||
return _seg_cache[model_name]
|
||||
|
||||
# 兼容旧默认 key
|
||||
if model_name == "mask2former_debug":
|
||||
model_name = "sam"
|
||||
|
||||
if model_name == "sam":
|
||||
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.SAM))
|
||||
_seg_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
if model_name == "mask2former":
|
||||
pred, _ = build_seg_predictor(UnifiedSegConfig(backend=SegBackend.MASK2FORMER))
|
||||
_seg_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 segment model_name: {model_name}")
|
||||
|
||||
|
||||
@app.post("/segment")
|
||||
def segment(req: SegmentRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "sam"
|
||||
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
rgb = np.array(pil, dtype=np.uint8)
|
||||
|
||||
predictor = _get_seg_predictor(model_name)
|
||||
label_map = predictor(rgb).astype(np.int32)
|
||||
|
||||
out_dir = OUTPUT_DIR / "segment"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = out_dir / f"{model_name}_label.png"
|
||||
# 保存为 8-bit 灰度(若 label 超过 255 会截断;当前 SAM 通常不会太大)
|
||||
Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)
|
||||
|
||||
return {"success": True, "label_path": str(out_path)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@app.post("/segment/sam_prompt")
|
||||
def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
|
||||
"""
|
||||
交互式 SAM:裁剪图 + 点/框提示,返回掩膜外轮廓点列(裁剪像素坐标)。
|
||||
"""
|
||||
try:
|
||||
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
rgb = np.array(pil, dtype=np.uint8)
|
||||
h, w = rgb.shape[0], rgb.shape[1]
|
||||
|
||||
if req.overlay_b64:
|
||||
ov = _b64_to_pil_image(req.overlay_b64)
|
||||
if ov.size != (w, h):
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
|
||||
"contour": [],
|
||||
}
|
||||
|
||||
if len(req.point_coords) != len(req.point_labels):
|
||||
return {
|
||||
"success": False,
|
||||
"error": "point_coords 与 point_labels 长度不一致",
|
||||
"contour": [],
|
||||
}
|
||||
if len(req.point_coords) < 1:
|
||||
return {"success": False, "error": "至少需要一个提示点", "contour": []}
|
||||
|
||||
pc = np.array(req.point_coords, dtype=np.float32)
|
||||
if pc.ndim != 2 or pc.shape[1] != 2:
|
||||
return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}
|
||||
|
||||
pl = np.array(req.point_labels, dtype=np.int64)
|
||||
box = np.array(req.box_xyxy, dtype=np.float32)
|
||||
|
||||
mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
|
||||
if not np.any(mask):
|
||||
return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}
|
||||
|
||||
contour = mask_to_contour_xy(mask, epsilon_px=2.0)
|
||||
if len(contour) < 3:
|
||||
return {"success": False, "error": "轮廓点数不足", "contour": []}
|
||||
|
||||
return {"success": True, "contour": contour, "error": None}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e), "contour": []}
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Inpaint
|
||||
# -----------------------------
|
||||
|
||||
|
||||
_inpaint_cache: Dict[str, Any] = {}
|
||||
|
||||
|
||||
def _get_inpaint_predictor(model_name: str):
|
||||
if model_name in _inpaint_cache:
|
||||
return _inpaint_cache[model_name]
|
||||
|
||||
if model_name == "copy":
|
||||
def _copy(image: Image.Image, *_args, **_kwargs) -> Image.Image:
|
||||
return image.convert("RGB")
|
||||
_inpaint_cache[model_name] = _copy
|
||||
return _copy
|
||||
|
||||
if model_name == "sdxl_inpaint":
|
||||
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.SDXL_INPAINT))
|
||||
_inpaint_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
if model_name == "controlnet":
|
||||
pred, _ = build_inpaint_predictor(UnifiedInpaintConfig(backend=InpaintBackend.CONTROLNET))
|
||||
_inpaint_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 inpaint model_name: {model_name}")
|
||||
|
||||
|
||||
@app.post("/inpaint")
|
||||
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "sdxl_inpaint"
|
||||
pil = _b64_to_pil_image(req.image_b64).convert("RGB")
|
||||
|
||||
if req.mask_b64:
|
||||
mask = _b64_to_pil_image(req.mask_b64).convert("L")
|
||||
else:
|
||||
mask = _default_half_mask(pil)
|
||||
|
||||
predictor = _get_inpaint_predictor(model_name)
|
||||
out = predictor(
|
||||
pil,
|
||||
mask,
|
||||
req.prompt or "",
|
||||
req.negative_prompt or "",
|
||||
strength=req.strength,
|
||||
max_side=req.max_side,
|
||||
)
|
||||
|
||||
out_dir = OUTPUT_DIR / "inpaint"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = out_dir / f"{model_name}_inpaint.png"
|
||||
out.save(out_path)
|
||||
|
||||
# 兼容 Qt 前端:直接返回结果图,避免前端再去读取服务器磁盘路径
|
||||
return {
|
||||
"success": True,
|
||||
"output_path": str(out_path),
|
||||
"output_image_b64": _pil_image_to_png_b64(out),
|
||||
}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
_animation_cache: Dict[str, Any] = {}
|
||||
|
||||
|
||||
def _get_animation_predictor(model_name: str):
|
||||
if model_name in _animation_cache:
|
||||
return _animation_cache[model_name]
|
||||
|
||||
if model_name == "animatediff":
|
||||
pred, _ = build_animation_predictor(
|
||||
UnifiedAnimationConfig(backend=AnimationBackend.ANIMATEDIFF)
|
||||
)
|
||||
_animation_cache[model_name] = pred
|
||||
return pred
|
||||
|
||||
raise ValueError(f"未知 animation model_name: {model_name}")
|
||||
|
||||
|
||||
@app.post("/animate")
|
||||
def animate(req: AnimateRequest) -> Dict[str, Any]:
|
||||
try:
|
||||
model_name = req.model_name or "animatediff"
|
||||
predictor = _get_animation_predictor(model_name)
|
||||
|
||||
result_path = predictor(
|
||||
prompt=req.prompt,
|
||||
negative_prompt=req.negative_prompt or "",
|
||||
num_inference_steps=req.num_inference_steps,
|
||||
guidance_scale=req.guidance_scale,
|
||||
width=req.width,
|
||||
height=req.height,
|
||||
video_length=req.video_length,
|
||||
seed=req.seed,
|
||||
)
|
||||
|
||||
out_dir = OUTPUT_DIR / "animation"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
out_path = out_dir / f"{model_name}_{ts}.gif"
|
||||
shutil.copy2(result_path, out_path)
|
||||
|
||||
return {"success": True, "output_path": str(out_path)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> Dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
from app.main import app
|
||||
|
||||
__all__ = ["app"]
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
port = int(os.environ.get("MODEL_SERVER_PORT", "8000"))
|
||||
uvicorn.run(app, host="0.0.0.0", port=port)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user