aa
This commit is contained in:
@@ -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 nose–L eye
|
||||
( 0, 230, 230), # 0→2 nose–R eye
|
||||
( 0, 180, 180), # 1→3 L eye–ear
|
||||
( 0, 180, 180), # 2→4 R eye–ear
|
||||
( 80, 80, 255), # 0→5 head–L shoulder
|
||||
(255, 80, 80), # 0→6 head–R 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"),
|
||||
|
||||
Reference in New Issue
Block a user