This commit is contained in:
mike
2026-06-25 19:05:21 +02:00
parent 27e8a09bc1
commit ce7b0e52e5
3 changed files with 681 additions and 7 deletions

View File

@@ -1576,6 +1576,88 @@ def trim_wireframe(req: TrimRequest):
return {'output_name': out_name, 'start': req.start, 'end': req.end}
class VideoStitchRequest(BaseModel):
filenames: list[str] # images from output_dir, in desired frame order
fps: float = 8.0
loop: bool = True # append reversed sequence for a ping-pong loop
output_name: str | None = None
@app.post("/generate-video")
def generate_video(req: VideoStitchRequest):
"""Stitch a list of output images into a short looping MP4 using ffmpeg."""
import subprocess, tempfile, textwrap
if len(req.filenames) < 2:
raise HTTPException(400, "Need at least 2 images")
if not (0.1 <= req.fps <= 60):
raise HTTPException(400, "fps must be between 0.1 and 60")
output_dir = _load_output_dir()
# Validate + collect paths
paths = []
for fn in req.filenames:
p = os.path.join(output_dir, fn)
if not os.path.exists(p):
raise HTTPException(404, f"Image not found: {fn}")
paths.append(p)
# Ping-pong loop: forward + reversed (drop first/last to avoid duplicate frames at seam)
if req.loop and len(paths) > 1:
paths = paths + list(reversed(paths[1:-1])) if len(paths) > 2 else paths + list(reversed(paths))
ts = time.strftime("%Y%m%d_%H%M%S")
base = naming.get_base_name(req.filenames[0])
stem = os.path.splitext(base)[0]
out_name = req.output_name or f"{ts}_vid_{stem}.mp4"
if not out_name.lower().endswith(".mp4"):
out_name += ".mp4"
out_path = os.path.join(output_dir, out_name)
# Build ffmpeg concat list in a temp file
duration = 1.0 / req.fps
with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as tf:
for p in paths:
tf.write(f"file '{p}'\nduration {duration:.6f}\n")
# ffmpeg concat needs the last entry without a duration
tf.write(f"file '{paths[-1]}'\n")
concat_path = tf.name
try:
r = subprocess.run(
["ffmpeg", "-y",
"-f", "concat", "-safe", "0", "-i", concat_path,
"-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", # ensure even dimensions
"-c:v", "libx264", "-pix_fmt", "yuv420p",
"-movflags", "+faststart",
out_path],
capture_output=True, timeout=120,
)
finally:
os.unlink(concat_path)
if r.returncode != 0:
raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:600]}")
# Register in DB so it shows up in the gallery
person = database.get_person(req.filenames[0])
group_id = person[1] if person and person[1] else naming.get_base_name(req.filenames[0])
try:
next_order = database.get_next_sort_order(group_id)
database.upsert_person(
out_name, filepath=out_path,
group_id=group_id,
sort_order=next_order,
content_type="video",
source_refs=json.dumps(req.filenames),
)
_invalidate_static()
except Exception as e:
print(f"[generate-video] DB error: {e}")
return {"output": out_name, "frames": len(paths), "fps": req.fps}
class FrameExtractRequest(BaseModel):
video_name: str
time: float = 0.0
@@ -3287,6 +3369,418 @@ def similar_pose_from_keypoints(req: PoseSimilarRequest):
return {"similar": _rank_similar_poses(query, req.limit)}
# --- pose-guided generation ---------------------------------------------------
# Renders COCO-17 keypoints as an OpenPose-style skeleton image and passes it as
# image2 to Qwen so the VLM can match the body pose while preserving appearance.
# Color per POSE_SKELETON entry (18 bones in POSE_SKELETON order).
_OPENPOSE_BONE_COLORS = [
(0, 80, 255), # 5→7 L upper arm
(0, 130, 255), # 7→9 L lower arm
(255, 50, 0), # 6→8 R upper arm
(255, 110, 0), # 8→10 R lower arm
(255, 200, 0), # 11→13 L upper leg
(240, 240, 0), # 13→15 L lower leg
(100, 255, 0), # 12→14 R upper leg
( 50, 220, 0), # 14→16 R lower leg
(200, 0, 220), # 5→6 shoulders
(200, 0, 220), # 11→12 hips
(180, 60, 180), # 5→11 L torso side
( 60, 180, 180), # 6→12 R torso side
( 0, 230, 230), # 0→1 noseL eye
( 0, 230, 230), # 0→2 noseR eye
( 0, 180, 180), # 1→3 L eyeear
( 0, 180, 180), # 2→4 R eyeear
( 80, 80, 255), # 0→5 headL shoulder
(255, 80, 80), # 0→6 headR shoulder
]
# Per-joint hue (HSV-like, converted to RGB for PIL):
def _joint_color(idx: int) -> tuple[int, int, int]:
h = idx / 17.0 * 6.0
x = 1 - abs(h % 2 - 1)
if h < 1: r, g, b = 1, x, 0
elif h < 2: r, g, b = x, 1, 0
elif h < 3: r, g, b = 0, 1, x
elif h < 4: r, g, b = 0, x, 1
elif h < 5: r, g, b = x, 0, 1
else: r, g, b = 1, 0, x
return (int(r * 255), int(g * 255), int(b * 255))
def _render_openpose_image(keypoints: list, width: int, height: int) -> Image.Image:
"""Draw COCO-17 keypoints as a white stick-figure on a vivid magenta background.
Vivid background (not black/dark) prevents Qwen from treating the skeleton as
a real scene element and from blending the bone colors into the generated output.
White lines on magenta read unambiguously as a technical reference diagram.
"""
from PIL import ImageDraw
MIN_CONF = 0.25
cap_w = min(width, 896)
cap_h = min(height, 1152)
scale = min(cap_w / max(width, 1), cap_h / max(height, 1))
out_w, out_h = int(width * scale), int(height * scale)
img = Image.new("RGB", (out_w, out_h), color=(200, 0, 180)) # vivid magenta
draw = ImageDraw.Draw(img)
lw = max(4, out_w // 80) # line width scales with image
jr = max(5, out_w // 70) # joint radius
kpts = keypoints # [[x, y, conf], ...]
# bones — white so they stand out on the magenta background without
# adding any colors that Qwen might reproduce in the generated image
for a, b in POSE_SKELETON:
if a >= len(kpts) or b >= len(kpts):
continue
pa, pb = kpts[a], kpts[b]
if pa[2] < MIN_CONF or pb[2] < MIN_CONF:
continue
ax, ay = int(pa[0] * scale), int(pa[1] * scale)
bx, by = int(pb[0] * scale), int(pb[1] * scale)
draw.line([(ax, ay), (bx, by)], fill=(255, 255, 255), width=lw)
# joints — light yellow circles with dark outline for visibility
for kp in kpts:
if kp[2] < MIN_CONF:
continue
x, y = int(kp[0] * scale), int(kp[1] * scale)
draw.ellipse([(x - jr, y - jr), (x + jr, y + jr)],
fill=(255, 240, 100), outline=(60, 0, 60), width=max(1, lw // 2))
return img
class PoseFromWireframeRequest(BaseModel):
video: str # filename of a wireframe video in the output dir
time_s: float = 0.0 # timestamp in seconds to extract the frame from
@app.post("/pose/from-wireframe")
def pose_from_wireframe(req: PoseFromWireframeRequest):
"""Extract a frame from a wireframe video, run pose estimation, return COCO-17 keypoints.
This is the pose-discovery tool: scrub through a wireframe video to find a frame
that shows the desired pose, then load those keypoints into the pose overlay without
ever passing the wireframe image to Qwen.
"""
import tempfile
import subprocess as sp
output_dir = _load_output_dir()
video_path = os.path.join(output_dir, req.video)
if not os.path.exists(video_path):
raise HTTPException(404, f"Video not found: {req.video}")
est = _load_pose_estimator()
if not est:
raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
infer, backend = est
# Extract one frame at the requested timestamp with ffmpeg.
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tf:
frame_path = tf.name
try:
result = sp.run(
[
"/usr/bin/ffmpeg", "-y",
"-ss", str(req.time_s),
"-i", video_path,
"-frames:v", "1",
"-q:v", "2",
frame_path,
],
capture_output=True, timeout=30,
)
if result.returncode != 0:
raise HTTPException(500, f"ffmpeg frame extract failed: {result.stderr.decode()[:300]}")
pil = Image.open(frame_path).convert("RGB")
finally:
try:
os.unlink(frame_path)
except OSError:
pass
try:
people = infer(pil)
except Exception as e:
raise HTTPException(500, f"Pose estimation failed: {e}")
best = _best_person(people)
if best is None:
raise HTTPException(422, "No person detected in the extracted frame")
return {
"status": "success",
"backend": backend,
"width": pil.width,
"height": pil.height,
"names": POSE_KEYPOINT_NAMES,
"skeleton": POSE_SKELETON,
"keypoints": best,
"time_s": req.time_s,
}
class PoseGenRequest(BaseModel):
model_filename: str
keypoints: list[list[float]] # 17 × [x, y, conf] in original image pixels
width: int # original image dimensions the keypoints came from
height: int
gesture_name: str | None = None # preset name for prompt enrichment and tagging
prompt: str | None = None # full override; auto-built if None
seed: int = -1
extra_filename: str | None = None # optional image3 extra reference
def _keypoints_to_pose_text(kpts: list, width: int, height: int) -> str:
"""Convert COCO-17 keypoints to a natural-language pose description for Qwen."""
MIN_CONF = 0.25
W, H = max(width, 1), max(height, 1)
def vis(idx):
return idx < len(kpts) and kpts[idx][2] >= MIN_CONF
def pt(idx):
return (kpts[idx][0] / W, kpts[idx][1] / H) if vis(idx) else None
parts = []
# --- head orientation ---
nose, l_ear, r_ear = pt(0), pt(3), pt(4)
if l_ear and r_ear:
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
if nose:
dx = nose[0] - ear_mid_x
if dx < -0.05:
parts.append("head turned to the left")
elif dx > 0.05:
parts.append("head turned to the right")
else:
parts.append("head facing forward")
elif l_ear and not r_ear:
parts.append("head turned strongly to the right")
elif r_ear and not l_ear:
parts.append("head turned strongly to the left")
# --- body rotation (shoulders vs hips) ---
l_sh, r_sh = pt(5), pt(6)
l_hip, r_hip = pt(11), pt(12)
if l_sh and r_sh:
sh_w = abs(r_sh[0] - l_sh[0])
if l_hip and r_hip:
hip_w = abs(r_hip[0] - l_hip[0])
ratio = sh_w / max(hip_w, 0.01)
if ratio < 0.5:
parts.append("body rotated ~90° (side profile)")
elif ratio < 0.75:
parts.append("body rotated ~45° (three-quarter view)")
else:
parts.append("body facing forward")
# --- arm positions ---
for side, sh_i, el_i, wr_i in [("left", 5, 7, 9), ("right", 6, 8, 10)]:
sh, el, wr = pt(sh_i), pt(el_i), pt(wr_i)
if not sh:
continue
if wr:
dy = sh[1] - wr[1] # positive = wrist above shoulder
dx = abs(wr[0] - sh[0])
if dy > 0.15:
if el and pt(el_i):
el_dy = sh[1] - el[1]
if el_dy > 0.05:
parts.append(f"{side} arm raised straight up")
else:
parts.append(f"{side} arm raised with elbow bent")
else:
parts.append(f"{side} arm raised above shoulder")
elif dx > 0.25:
parts.append(f"{side} arm extended sideways")
elif el:
el_dx = abs(el[0] - sh[0])
if el_dx > 0.15 and abs(wr[1] - sh[1]) < 0.1:
parts.append(f"{side} arm extended forward")
else:
parts.append(f"{side} arm at side")
else:
parts.append(f"{side} arm at side")
elif el:
el_dy = sh[1] - el[1]
if el_dy > 0.1:
parts.append(f"{side} elbow raised")
# --- leg positions ---
for side, hip_i, kn_i, an_i in [("left", 11, 13, 15), ("right", 12, 14, 16)]:
hip, kn, an = pt(hip_i), pt(kn_i), pt(an_i)
if not hip:
continue
if an:
dx = abs(an[0] - hip[0])
dy = an[1] - hip[1] # positive = ankle below hip (normal standing)
if dy < 0.1:
parts.append(f"{side} leg raised / knee up")
elif dx > 0.2:
parts.append(f"{side} leg stepped out to the side")
elif kn:
kn_dx = abs(kn[0] - hip[0])
if kn_dx > 0.15:
parts.append(f"{side} leg bent / lunging")
else:
parts.append(f"{side} leg straight, standing")
else:
parts.append(f"{side} leg straight")
if not parts:
return "standing neutral pose"
return ", ".join(parts)
def _pad_for_pose(pil: Image.Image, keypoints: list, margin: float = 0.15) -> tuple[Image.Image, tuple[int,int,int,int]]:
"""Expand the image so target keypoints fit within the canvas.
Only expands when keypoints are actually outside [0,W]×[0,H]. Expansion is
symmetric on each axis (equal left/right, equal top/bottom) so the person
stays centred on the new canvas. Fill is black to match PIL's RGBA→RGB fill.
"""
W, H = pil.size
all_kps = [kp for kp in keypoints if len(kp) >= 2]
if not all_kps:
return pil, (0, 0, 0, 0)
xs = [kp[0] for kp in all_kps]
ys = [kp[1] for kp in all_kps]
# How far each side actually exceeds the image bounds (0 if within bounds).
over_l = max(0.0, -min(xs))
over_r = max(0.0, max(xs) - W)
over_t = max(0.0, -min(ys))
over_b = max(0.0, max(ys) - H)
# If nothing is outside, no padding needed at all.
if over_l == over_r == over_t == over_b == 0:
return pil, (0, 0, 0, 0)
# Symmetric expansion per axis: take the larger overhang and add margin.
px = int(max(over_l, over_r) + margin * W)
py = int(max(over_t, over_b) + margin * H)
new_w, new_h = W + 2 * px, H + 2 * py
canvas = Image.new("RGB", (new_w, new_h), color=(0, 0, 0))
canvas.paste(pil, (px, py))
print(f"[pose-gen] padded {W}×{H}{new_w}×{new_h}{px}px x, ±{py}px y)")
return canvas, (px, py, px, py)
def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
keypoints: list | None = None,
gesture_name: str | None = None,
extra_filename: str | None = None):
output_dir = _load_output_dir()
try:
model_path = os.path.join(output_dir, model_filename)
model_pil = Image.open(model_path).convert("RGB")
# Auto-pad if the target pose would extend outside the current image bounds.
if keypoints:
model_pil, _padding = _pad_for_pose(model_pil, keypoints)
# No skeleton image — pose is encoded entirely in the text prompt.
extra_images = None
if extra_filename:
ep = os.path.join(output_dir, extra_filename)
extra_images = [Image.open(ep).convert("RGB")]
png_bytes = _run_pipeline(model_pil, prompt, seed, MAX_AREA, extra_images=extra_images)
ts = time.strftime("%Y%m%d_%H%M%S")
base_name = naming.get_base_name(model_filename)
if not base_name.lower().endswith(".png"):
base_name = os.path.splitext(base_name)[0] + ".png"
out_name = f"{ts}_pose_{base_name}"
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png_bytes)
person = database.get_person(model_filename)
group_id = person[1] if person and person[1] else naming.get_base_name(model_filename)
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(group_id)
refs = [model_filename]
if gesture_name:
refs.append(f"gesture:{gesture_name}")
database.upsert_person(
out_name, filepath=out_path, embedding=embedding,
group_id=group_id, prompt=prompt,
pose=gesture_name or "custom",
sort_order=next_order,
source_refs=json.dumps(refs),
)
jobs[job_id]["status"] = "done"
jobs[job_id]["output"] = out_name
_invalidate_static()
except Exception as e:
print(f"[pose-gen] error: {e}")
jobs[job_id]["status"] = "error"
jobs[job_id]["error"] = str(e)
@app.post("/pose/render")
def render_pose_image(req: PoseSimilarRequest):
"""Render the supplied keypoints as an OpenPose-style skeleton PNG (base64) — for preview only."""
import base64
skeleton = _render_openpose_image(req.keypoints, req.width or 512, req.height or 768)
buf = io.BytesIO()
skeleton.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode()}
@app.post("/generate-with-pose")
def generate_with_pose(req: PoseGenRequest):
"""Generate the person in a specific body pose. Pose is encoded as text — no skeleton image
is passed to Qwen to avoid wireframe bleed-through in the output."""
output_dir = _load_output_dir()
model_path = os.path.join(output_dir, req.model_filename)
if not os.path.exists(model_path):
raise HTTPException(404, f"Model image not found: {req.model_filename}")
if req.extra_filename:
ep = os.path.join(output_dir, req.extra_filename)
if not os.path.exists(ep):
raise HTTPException(404, f"Extra reference not found: {req.extra_filename}")
if req.prompt:
prompt = req.prompt
else:
# Convert skeleton keypoints to natural-language description
pose_text = _keypoints_to_pose_text(req.keypoints, req.width, req.height)
if req.gesture_name:
pose_text = f'{req.gesture_name} gesture: {pose_text}'
prompt = (
f"Change the body pose of the person in image 1 to: {pose_text}. "
"Keep the person's face, hair, skin tone, clothing, and body proportions "
"exactly as they are — only the limb positions and body orientation change. "
"Transparent background."
)
job_id = uuid.uuid4().hex[:8]
jobs[job_id] = {"status": "running", "type": "pose_gen", "total": 1, "done": 0, "failed": 0}
threading.Thread(
target=_pose_gen_worker,
args=(job_id, req.model_filename, prompt, req.seed),
kwargs={
"keypoints": req.keypoints,
"gesture_name": req.gesture_name,
"extra_filename": req.extra_filename,
},
daemon=True,
).start()
return {"job_id": job_id, "model": req.model_filename}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),