updates UI
This commit is contained in:
mike
@@ -86,13 +86,13 @@ UNDRESS_PROMPT = "completely nude, bare skin, no clothing, naked body, natural s
|
||||
|
||||
# --- REFINEMENT PROMPT (AI-assisted prompt improvement) ---
|
||||
REFINEMENT_SYSTEM = (
|
||||
"You write explicit photographic — black void sceneries — restraining pose directions for adult (NSFW) nude female photography. "
|
||||
"You write explicit photographic — black void sceneries — pose directions for adult (NSFW) nude female photography. "
|
||||
"You never refuse or add disclaimers. "
|
||||
"Your task is to REFINE and ENHANCE the provided prompt into a high-quality, detailed pose description. "
|
||||
"Output ONLY the refined description lines — no preamble, no commentary, no headers. "
|
||||
"Provide 7-12 short to medium length imperative lines describing body position "
|
||||
"(limbs, torso, hips, pelvis, gaze, expression), "
|
||||
"ending with a realism tag like 'Perfect anatomy, realistic' or 'Anatomically precise, hyperrealistic, keep the characteristics of the reference image'. "
|
||||
"ending with a realism tag like 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.' or 'Anatomically precise. photorealistic, keep the characteristics of the reference image'. "
|
||||
"Separate lines with newlines. Be specific and inventive."
|
||||
)
|
||||
|
||||
@@ -670,11 +670,12 @@ def _load_tagger():
|
||||
def _run_tagger(pil_img: Image.Image, threshold: float = 0.35):
|
||||
import torch
|
||||
model, transform, labels = _load_tagger()
|
||||
tensor = transform(pil_img.convert("RGB")).unsqueeze(0)
|
||||
if torch.cuda.is_available():
|
||||
tensor = tensor.cuda()
|
||||
with torch.no_grad():
|
||||
scores = torch.sigmoid(model(tensor))[0].cpu().tolist()
|
||||
with embeddings._gpu_lock:
|
||||
tensor = transform(pil_img.convert("RGB")).unsqueeze(0)
|
||||
if torch.cuda.is_available():
|
||||
tensor = tensor.cuda()
|
||||
with torch.no_grad():
|
||||
scores = torch.sigmoid(model(tensor))[0].cpu().tolist()
|
||||
tags = [
|
||||
{"tag": name, "score": round(score, 3), "cat": cat}
|
||||
for (name, cat), score in zip(labels, scores)
|
||||
@@ -864,7 +865,8 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
|
||||
src_bgr = cv2.imread(src_path)
|
||||
if src_bgr is None:
|
||||
raise ValueError(f"Cannot read model image: {model_filename}")
|
||||
src_faces = app.get(src_bgr)
|
||||
with embeddings._gpu_lock:
|
||||
src_faces = app.get(src_bgr)
|
||||
if not src_faces:
|
||||
raise ValueError(f"No face detected in: {model_filename}")
|
||||
# Use the largest face as source
|
||||
@@ -912,7 +914,8 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
tgt_faces = app.get(frame)
|
||||
with embeddings._gpu_lock:
|
||||
tgt_faces = app.get(frame)
|
||||
result = frame
|
||||
if tgt_faces:
|
||||
# Only swap the largest face — avoids false-positive detections
|
||||
@@ -925,14 +928,16 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
|
||||
result = frame.copy()
|
||||
best_face = max(valid, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||||
try:
|
||||
result = swapper.get(result, best_face, src_face, paste_back=True)
|
||||
with embeddings._gpu_lock:
|
||||
result = swapper.get(result, best_face, src_face, paste_back=True)
|
||||
except Exception:
|
||||
pass
|
||||
if gfpgan_restorer is not None:
|
||||
try:
|
||||
_, _, result = gfpgan_restorer.enhance(
|
||||
result, has_aligned=False, only_center_face=False, paste_back=True
|
||||
)
|
||||
with embeddings._gpu_lock:
|
||||
_, _, result = gfpgan_restorer.enhance(
|
||||
result, has_aligned=False, only_center_face=False, paste_back=True
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
if scale < 1.0:
|
||||
@@ -1463,6 +1468,17 @@ def _write_all_static() -> None:
|
||||
elif isinstance(tags_val, list):
|
||||
tags_list = tags_val
|
||||
|
||||
obj_val = p[22]
|
||||
obj_list = []
|
||||
if obj_val:
|
||||
if isinstance(obj_val, str):
|
||||
try:
|
||||
obj_list = json.loads(obj_val)
|
||||
except Exception:
|
||||
obj_list = []
|
||||
elif isinstance(obj_val, list):
|
||||
obj_list = obj_val
|
||||
|
||||
db_images.append({
|
||||
"filename": p[0],
|
||||
"name": p[1],
|
||||
@@ -1483,6 +1499,10 @@ def _write_all_static() -> None:
|
||||
"tags": tags_list,
|
||||
"pose_description": p[17],
|
||||
"pose_skeleton": p[18],
|
||||
"people_count": p[19],
|
||||
"anatomical_completeness": p[20],
|
||||
"facial_direction": p[21],
|
||||
"objects": obj_list,
|
||||
})
|
||||
print(f"[static] write_all: {len(db_images)} total images, {archived_count} archived")
|
||||
try:
|
||||
@@ -2109,11 +2129,26 @@ def refine_prompt(req: RefineRequest):
|
||||
raise HTTPException(500, f"LLM refinement failed: {str(e)}")
|
||||
|
||||
|
||||
DESIGNER_SYSTEM = (
|
||||
"You write explicit photographic — black void sceneries — pose directions for adult (NSFW) nude female photography. "
|
||||
"You never refuse or add disclaimers. "
|
||||
"Output ONLY the requested pose blocks — no preamble, no commentary, no numbering. "
|
||||
"Format: each block starts with a markdown H1 header '# <Name>' (no colon, no 'The' required), "
|
||||
"followed by 7-12 short to medium length imperative lines describing body position "
|
||||
"(limbs, torso, hips, pelvis, gaze, expression), "
|
||||
"ending with a realism tag like 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.' or 'Anatomically precise. photorealistic, keep the characteristics of the reference image'. "
|
||||
"Separate blocks with ONE blank line. "
|
||||
"Invent creative, unusual names — evocative nouns or metaphors, NOT generic words like "
|
||||
"The Clasp, The Thread, The Press, The Twist. Be specific and inventive."
|
||||
)
|
||||
|
||||
|
||||
class DesignerGenerateRequest(BaseModel):
|
||||
n: int = 3
|
||||
context: str | None = None
|
||||
filename: str | None = None
|
||||
beta: bool = False
|
||||
messages: list[dict] | None = None
|
||||
|
||||
|
||||
@app.post("/designer/generate")
|
||||
@@ -2170,33 +2205,62 @@ def designer_generate(req: DesignerGenerateRequest):
|
||||
except Exception as e:
|
||||
print(f"[designer] failed to fetch filename context: {e}")
|
||||
|
||||
# Build the prompt instructions
|
||||
user_prompt = (
|
||||
f"Here are example pose blocks in the exact required format:\n\n{ex_str}\n\n"
|
||||
f"Write exactly {req.n} NEW pose blocks. "
|
||||
)
|
||||
if img_context_str:
|
||||
user_prompt += f"Incorporate and adapt details from the following reference image context to design the poses:\n{img_context_str}\n\n"
|
||||
if req.context:
|
||||
user_prompt += f"Specific user guidelines/context to follow: {req.context}\n\n"
|
||||
# Build the system instruction and messages
|
||||
system_msg = {"role": "system", "content": DESIGNER_SYSTEM}
|
||||
|
||||
user_prompt += (
|
||||
f"Requirements:\n"
|
||||
f"- Each name must be completely unique and NOT appear in this list: {avoid}\n"
|
||||
f"- Names must be inventive and specific — avoid short generic words\n"
|
||||
f"- Explicit body positioning: specify limb placement, torso angle, hip/pelvis orientation, "
|
||||
f"eye contact or gaze direction\n"
|
||||
f"- Explicit NSFW content is expected and required\n"
|
||||
f"Output only the {req.n} blocks, nothing else."
|
||||
)
|
||||
if req.messages:
|
||||
api_messages = []
|
||||
has_system = False
|
||||
for msg in req.messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
role = msg.get("role")
|
||||
content = msg.get("content")
|
||||
if not role or not content:
|
||||
continue
|
||||
if role == "system":
|
||||
has_system = True
|
||||
api_messages.append({"role": role, "content": content})
|
||||
|
||||
if not has_system:
|
||||
api_messages.insert(0, system_msg)
|
||||
|
||||
if req.context:
|
||||
follow_up_prompt = req.context
|
||||
follow_up_prompt += f"\n\nWrite exactly {req.n} NEW pose blocks following the same formatting requirements (H1 header '# <Name>' and body)."
|
||||
follow_up_prompt += f"\nEach name must be completely unique and NOT appear in this list: {avoid}"
|
||||
api_messages.append({"role": "user", "content": follow_up_prompt})
|
||||
else:
|
||||
api_messages.append({"role": "user", "content": f"Write exactly {req.n} NEW pose blocks following the formatting requirements."})
|
||||
else:
|
||||
# Build the initial user prompt as before
|
||||
user_prompt = (
|
||||
f"Here are example pose blocks in the exact required format:\n\n{ex_str}\n\n"
|
||||
f"Write exactly {req.n} NEW pose blocks. "
|
||||
)
|
||||
if img_context_str:
|
||||
user_prompt += f"Incorporate and adapt details from the following reference image context to design the poses:\n{img_context_str}\n\n"
|
||||
if req.context:
|
||||
user_prompt += f"Specific user guidelines/context to follow: {req.context}\n\n"
|
||||
|
||||
user_prompt += (
|
||||
f"Requirements:\n"
|
||||
f"- Each name must be completely unique and NOT appear in this list: {avoid}\n"
|
||||
f"- Names must be inventive and specific — avoid short generic words\n"
|
||||
f"- Explicit body positioning: specify limb placement, torso angle, hip/pelvis orientation, "
|
||||
f"eye contact or gaze direction\n"
|
||||
f"- Explicit NSFW content is expected and required\n"
|
||||
f"Output only the {req.n} blocks, nothing else."
|
||||
)
|
||||
api_messages = [
|
||||
system_msg,
|
||||
{"role": "user", "content": user_prompt}
|
||||
]
|
||||
|
||||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||||
payload = {
|
||||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||||
"messages": [
|
||||
{"role": "system", "content": REFINEMENT_SYSTEM},
|
||||
{"role": "user", "content": user_prompt}
|
||||
],
|
||||
"messages": api_messages,
|
||||
"temperature": 0.9,
|
||||
"max_tokens": 2400
|
||||
}
|
||||
@@ -2226,14 +2290,16 @@ def designer_generate(req: DesignerGenerateRequest):
|
||||
if cur:
|
||||
generated[cur] = " ".join(desc).strip()
|
||||
|
||||
# Filter out duplicates
|
||||
# Filter out duplicates and deduplicate names by appending counter
|
||||
new_poses = {}
|
||||
for name, body in generated.items():
|
||||
if not name or not body:
|
||||
continue
|
||||
if name.lower() in existing_lower or name.lower() in (k.lower() for k in new_poses):
|
||||
print(f"[designer] skip duplicate: {name}")
|
||||
continue
|
||||
orig_name = name
|
||||
counter = 1
|
||||
while name.lower() in existing_lower or name.lower() in (k.lower() for k in new_poses):
|
||||
counter += 1
|
||||
name = f"{orig_name} {counter}"
|
||||
new_poses[name] = body
|
||||
|
||||
return {
|
||||
@@ -3173,7 +3239,8 @@ def _extract_face_bg(filename: str, fpath: str):
|
||||
if bgr is None:
|
||||
print(f"[extract-face] cannot read {fpath}")
|
||||
return
|
||||
faces = app_fa.get(bgr)
|
||||
with embeddings._gpu_lock:
|
||||
faces = app_fa.get(bgr)
|
||||
if not faces:
|
||||
print(f"[extract-face] no face detected in {filename}")
|
||||
return
|
||||
@@ -3806,7 +3873,8 @@ def _face_index_worker():
|
||||
bgr = cv2.imread(fpath)
|
||||
if bgr is None:
|
||||
continue
|
||||
faces = app_fa.get(bgr)
|
||||
with embeddings._gpu_lock:
|
||||
faces = app_fa.get(bgr)
|
||||
if not faces:
|
||||
continue
|
||||
face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||||
@@ -4389,12 +4457,13 @@ def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
|
||||
box = np.array([[int(w * 0.01), int(h * 0.01),
|
||||
int(w * 0.99), int(h * 0.99)]], dtype=np.float32)
|
||||
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=box,
|
||||
multimask_output=True,
|
||||
)
|
||||
with embeddings._gpu_lock:
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=box,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
if masks is None or len(masks) == 0:
|
||||
print("[sam2] no masks returned, falling back to rembg")
|
||||
@@ -4500,9 +4569,10 @@ def _apply_transparency_black_bg(png_bytes: bytes) -> bytes:
|
||||
if predictor is not False:
|
||||
box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
||||
with embeddings._gpu_lock:
|
||||
with torch.inference_mode():
|
||||
predictor.set_image(arr)
|
||||
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
||||
|
||||
if masks is not None and len(masks) > 0:
|
||||
best = masks[int(np.argmax(scores))]
|
||||
@@ -5144,33 +5214,109 @@ def _pose_distance(a, b):
|
||||
return min(direct, mirror)
|
||||
|
||||
def _describe_pose(kpts):
|
||||
"""Generate a simple human-readable description of a COCO-17 pose."""
|
||||
"""Generate a highly detailed human-readable description of a COCO-17 pose,
|
||||
including body orientation (standing, sitting, laying down), facing direction,
|
||||
arms position, and legs posture.
|
||||
"""
|
||||
vis = [k[2] >= _POSE_MIN_SCORE for k in kpts]
|
||||
if sum(vis) < 5: return "Indeterminate pose"
|
||||
|
||||
parts = []
|
||||
|
||||
# Vertical orientation
|
||||
if vis[0] and vis[11] and vis[12]: # nose and hips
|
||||
hip_y = (kpts[11][1] + kpts[12][1]) / 2
|
||||
head_y = kpts[0][1]
|
||||
if head_y > hip_y + 20: parts.append("upside down")
|
||||
elif head_y > hip_y - 20: parts.append("reclining/prone")
|
||||
else: parts.append("upright")
|
||||
# Extract some key positions
|
||||
head_x = kpts[0][0] if vis[0] else ((kpts[5][0] + kpts[6][0])/2 if (vis[5] and vis[6]) else None)
|
||||
head_y = kpts[0][1] if vis[0] else ((kpts[5][1] + kpts[6][1])/2 if (vis[5] and vis[6]) else None)
|
||||
|
||||
# Arms
|
||||
if vis[9] and vis[10]: # wrists
|
||||
sh_y = (kpts[5][1] + kpts[6][1]) / 2 if (vis[5] and vis[6]) else kpts[0][1]
|
||||
if kpts[9][1] < sh_y and kpts[10][1] < sh_y: parts.append("arms raised")
|
||||
elif kpts[9][1] > sh_y + 100 and kpts[10][1] > sh_y + 100: parts.append("arms down")
|
||||
else: parts.append("arms at sides")
|
||||
|
||||
# Legs
|
||||
if vis[15] and vis[16]: # ankles
|
||||
dist = abs(kpts[15][0] - kpts[16][0])
|
||||
if dist > 150: parts.append("legs spread")
|
||||
else: parts.append("legs together")
|
||||
|
||||
hip_x = (kpts[11][0] + kpts[12][0])/2 if (vis[11] and vis[12]) else (kpts[11][0] if vis[11] else (kpts[12][0] if vis[12] else None))
|
||||
hip_y = (kpts[11][1] + kpts[12][1])/2 if (vis[11] and vis[12]) else (kpts[11][1] if vis[11] else (kpts[12][1] if vis[12] else None))
|
||||
|
||||
sh_y = (kpts[5][1] + kpts[6][1])/2 if (vis[5] and vis[6]) else (kpts[0][1] if vis[0] else None)
|
||||
torso_h = abs(hip_y - sh_y) if (hip_y is not None and sh_y is not None) else 100.0
|
||||
|
||||
# 1. Posture (standing, sitting, kneeling/crouching, laying down/reclining)
|
||||
posture = "upright"
|
||||
if head_y is not None and hip_y is not None:
|
||||
if head_y > hip_y + 30:
|
||||
posture = "upside down"
|
||||
else:
|
||||
# Check horizontal vs vertical distance to identify lying down
|
||||
dx = abs(head_x - hip_x) if head_x is not None and hip_x is not None else 0
|
||||
dy = abs(head_y - hip_y)
|
||||
if dx > 1.2 * dy:
|
||||
posture = "lying down/reclining"
|
||||
|
||||
if posture == "upright":
|
||||
has_hips = vis[11] or vis[12]
|
||||
has_knees = vis[13] or vis[14]
|
||||
has_ankles = vis[15] or vis[16]
|
||||
if has_hips and has_knees:
|
||||
h_y = hip_y
|
||||
h_x = hip_x
|
||||
k_y = (kpts[13][1] + kpts[14][1])/2 if (vis[13] and vis[14]) else (kpts[13][1] if vis[13] else kpts[14][1])
|
||||
k_x = (kpts[13][0] + kpts[14][0])/2 if (vis[13] and vis[14]) else (kpts[13][0] if vis[13] else kpts[14][0])
|
||||
|
||||
thigh_dy = abs(k_y - h_y)
|
||||
thigh_dx = abs(k_x - h_x)
|
||||
|
||||
if has_ankles:
|
||||
a_y = (kpts[15][1] + kpts[16][1])/2 if (vis[15] and vis[16]) else (kpts[15][1] if vis[15] else kpts[16][1])
|
||||
a_x = (kpts[15][0] + kpts[16][0])/2 if (vis[15] and vis[16]) else (kpts[15][0] if vis[15] else kpts[16][0])
|
||||
shin_dy = abs(a_y - k_y)
|
||||
shin_dx = abs(a_x - k_x)
|
||||
|
||||
# Sitting: thigh horizontal, shin vertical
|
||||
if thigh_dy < 0.6 * thigh_dx and shin_dy > 1.2 * shin_dx:
|
||||
posture = "sitting"
|
||||
elif thigh_dy < 0.45 * torso_h and shin_dy > 0.5 * torso_h:
|
||||
posture = "sitting"
|
||||
# Crouching/Kneeling: hips close to ankles/ground
|
||||
elif abs(h_y - a_y) < 0.85 * torso_h:
|
||||
posture = "crouching/kneeling"
|
||||
else:
|
||||
posture = "standing"
|
||||
else:
|
||||
# Ankles not visible
|
||||
if thigh_dy < 0.5 * torso_h:
|
||||
posture = "sitting"
|
||||
else:
|
||||
posture = "standing"
|
||||
|
||||
parts.append(posture)
|
||||
|
||||
# 2. Body Orientation / Facing Direction
|
||||
if vis[5] and vis[6]:
|
||||
sh_dist = abs(kpts[5][0] - kpts[6][0])
|
||||
# Profile view check (shoulders compressed horizontally)
|
||||
if sh_dist < 0.25 * torso_h:
|
||||
parts.append("turned sideways (profile view)")
|
||||
# Back view check
|
||||
elif kpts[5][0] < kpts[6][0]:
|
||||
parts.append("facing away (back view)")
|
||||
else:
|
||||
parts.append("facing forward (front view)")
|
||||
|
||||
# 3. Arms Posture
|
||||
if vis[9] and vis[10]: # wrists
|
||||
sh_y_val = sh_y if sh_y is not None else kpts[0][1]
|
||||
if kpts[9][1] < sh_y_val and kpts[10][1] < sh_y_val:
|
||||
parts.append("arms raised")
|
||||
elif kpts[9][1] > sh_y_val + torso_h * 0.8 and kpts[10][1] > sh_y_val + torso_h * 0.8:
|
||||
# Check if close to hips (hands on hips)
|
||||
if (vis[11] and abs(kpts[9][0] - kpts[11][0]) < torso_h * 0.2) or (vis[12] and abs(kpts[10][0] - kpts[12][0]) < torso_h * 0.2):
|
||||
parts.append("hands on hips")
|
||||
else:
|
||||
parts.append("arms down")
|
||||
else:
|
||||
parts.append("arms at sides")
|
||||
|
||||
# 4. Legs Posture
|
||||
if vis[15] and vis[16]: # ankles
|
||||
ankle_dist = abs(kpts[15][0] - kpts[16][0])
|
||||
if ankle_dist > torso_h * 0.6:
|
||||
parts.append("legs spread")
|
||||
else:
|
||||
parts.append("legs together")
|
||||
|
||||
if not parts: return "Generic pose"
|
||||
return ", ".join(parts)
|
||||
|
||||
@@ -6109,67 +6255,309 @@ def _detect_people_count(keypoints: list) -> int:
|
||||
return 1 if keypoints else 0
|
||||
|
||||
|
||||
def _detect_anatomical_completeness(keypoints: list) -> bool:
|
||||
def _detect_anatomical_completeness(keypoints: list, width: int = None, height: int = None) -> bool:
|
||||
"""Detect if the person has complete anatomical structure.
|
||||
|
||||
Returns True if all major body parts are visible (head, torso, arms, legs).
|
||||
Uses pose keypoint visibility to determine completeness.
|
||||
Returns True if all major body parts are visible (head, torso, arms, legs)
|
||||
and are fully contained within the frame (not cropped at boundaries).
|
||||
"""
|
||||
if not keypoints or len(keypoints) < 17:
|
||||
return False
|
||||
|
||||
# Minimum visibility threshold for each keypoint
|
||||
MIN_VISIBILITY = 0.3
|
||||
|
||||
# Key keypoints that indicate anatomical completeness
|
||||
# Head (0), shoulders (5,6), hips (11,12), elbows (7,8), wrists (9,10), knees (13,14), ankles (15,16)
|
||||
keypoint_indices = [0, 5, 6, 11, 12, 7, 8, 9, 10, 13, 14, 15, 16]
|
||||
# 1. Presence of major body segments (requires visibility >= 0.3)
|
||||
has_head = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [0, 1, 2, 3, 4])
|
||||
has_shoulders = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [5, 6])
|
||||
has_hips = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [11, 12])
|
||||
has_knees = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [13, 14])
|
||||
has_ankles = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [15, 16])
|
||||
|
||||
visible_count = 0
|
||||
for idx in keypoint_indices:
|
||||
if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY:
|
||||
visible_count += 1
|
||||
|
||||
# If more than half of the key keypoints are visible, consider it complete
|
||||
return visible_count > len(keypoint_indices) * 0.5
|
||||
# If any major body segment is completely missing, the anatomy is not complete (e.g., cropped above knees/hips/chest)
|
||||
if not (has_head and has_shoulders and has_hips and has_knees and has_ankles):
|
||||
return False
|
||||
|
||||
# 2. Boundary cropping check (if image dimensions are provided)
|
||||
if width and height:
|
||||
# Check if head is too close to top edge
|
||||
head_kpts = [keypoints[idx] for idx in [0, 1, 2, 3, 4] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||||
if head_kpts:
|
||||
min_head_y = min(kp[1] for kp in head_kpts)
|
||||
if min_head_y < height * 0.02: # head is cropped at top
|
||||
return False
|
||||
|
||||
# Check if ankles (feet) are too close to bottom edge
|
||||
ankle_kpts = [keypoints[idx] for idx in [15, 16] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||||
if ankle_kpts:
|
||||
max_ankle_y = max(kp[1] for kp in ankle_kpts)
|
||||
if max_ankle_y > height * 0.98: # ankles/feet are cropped at bottom
|
||||
return False
|
||||
|
||||
# Check if wrists (hands) are too close to left/right edge
|
||||
wrist_kpts = [keypoints[idx] for idx in [9, 10] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||||
if wrist_kpts:
|
||||
min_wrist_x = min(kp[0] for kp in wrist_kpts)
|
||||
max_wrist_x = max(kp[0] for kp in wrist_kpts)
|
||||
if min_wrist_x < width * 0.02 or max_wrist_x > width * 0.98: # wrists/hands are cropped at sides
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _detect_facial_direction(keypoints: list) -> str:
|
||||
"""Detect the facial direction from keypoints.
|
||||
"""Detect the facial direction from keypoints, including fine-grained gaze details (up/down/left/right).
|
||||
|
||||
Returns a string describing the head orientation.
|
||||
Returns a string describing the head/gaze orientation.
|
||||
"""
|
||||
if not keypoints or len(keypoints) < 17:
|
||||
return "unknown"
|
||||
|
||||
# Key points for face direction detection
|
||||
# Nose (0), left ear (3), right ear (4)
|
||||
# Nose (0), left eye (1), right eye (2), left ear (3), right ear (4)
|
||||
nose = keypoints[0] if len(keypoints) > 0 and keypoints[0][2] >= 0.3 else None
|
||||
l_eye = keypoints[1] if len(keypoints) > 1 and keypoints[1][2] >= 0.3 else None
|
||||
r_eye = keypoints[2] if len(keypoints) > 2 and keypoints[2][2] >= 0.3 else None
|
||||
l_ear = keypoints[3] if len(keypoints) > 3 and keypoints[3][2] >= 0.3 else None
|
||||
r_ear = keypoints[4] if len(keypoints) > 4 and keypoints[4][2] >= 0.3 else None
|
||||
|
||||
if not nose:
|
||||
return "unknown"
|
||||
|
||||
# Determine face direction based on ear positions
|
||||
# 1. Determine horizontal direction
|
||||
horiz = "forward"
|
||||
if l_ear and r_ear:
|
||||
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
|
||||
dx = nose[0] - ear_mid_x
|
||||
if dx < -0.05:
|
||||
return "looking left"
|
||||
elif dx > 0.05:
|
||||
return "looking right"
|
||||
else:
|
||||
return "looking forward"
|
||||
# Normalize dx by distance between ears to make it scale invariant
|
||||
ear_dist = abs(l_ear[0] - r_ear[0])
|
||||
if ear_dist > 0:
|
||||
norm_dx = dx / ear_dist
|
||||
if norm_dx < -0.06:
|
||||
horiz = "left"
|
||||
elif norm_dx > 0.06:
|
||||
horiz = "right"
|
||||
elif l_ear and not r_ear:
|
||||
return "looking strongly right"
|
||||
horiz = "strongly right"
|
||||
elif r_ear and not l_ear:
|
||||
return "looking strongly left"
|
||||
else:
|
||||
return "looking forward"
|
||||
horiz = "strongly left"
|
||||
|
||||
# 2. Determine vertical direction (up/down)
|
||||
vert = "level"
|
||||
# Try using eyes first as they are closer to the nose
|
||||
if l_eye and r_eye:
|
||||
eye_y = (l_eye[1] + r_eye[1]) / 2
|
||||
eye_dist = abs(l_eye[0] - r_eye[0])
|
||||
if eye_dist > 0:
|
||||
v_ratio = (nose[1] - eye_y) / eye_dist
|
||||
if v_ratio < 0.15:
|
||||
vert = "up"
|
||||
elif v_ratio > 0.65:
|
||||
vert = "down"
|
||||
elif l_ear and r_ear:
|
||||
ear_y = (l_ear[1] + r_ear[1]) / 2
|
||||
ear_dist = abs(l_ear[0] - r_ear[0])
|
||||
if ear_dist > 0:
|
||||
v_ratio = (nose[1] - ear_y) / ear_dist
|
||||
if v_ratio < -0.1:
|
||||
vert = "up"
|
||||
elif v_ratio > 0.3:
|
||||
vert = "down"
|
||||
|
||||
# 3. Combine horizontal and vertical
|
||||
if horiz == "forward":
|
||||
if vert == "level":
|
||||
return "looking forward"
|
||||
elif vert == "up":
|
||||
return "looking forward and up"
|
||||
elif vert == "down":
|
||||
return "looking forward and down"
|
||||
elif horiz == "left":
|
||||
if vert == "level":
|
||||
return "looking left"
|
||||
elif vert == "up":
|
||||
return "looking left and up"
|
||||
elif vert == "down":
|
||||
return "looking left and down"
|
||||
elif horiz == "right":
|
||||
if vert == "level":
|
||||
return "looking right"
|
||||
elif vert == "up":
|
||||
return "looking right and up"
|
||||
elif vert == "down":
|
||||
return "looking right and down"
|
||||
elif horiz == "strongly left":
|
||||
if vert == "level":
|
||||
return "looking strongly left"
|
||||
elif vert == "up":
|
||||
return "looking strongly left and up"
|
||||
elif vert == "down":
|
||||
return "looking strongly left and down"
|
||||
elif horiz == "strongly right":
|
||||
if vert == "level":
|
||||
return "looking strongly right"
|
||||
elif vert == "up":
|
||||
return "looking strongly right and up"
|
||||
elif vert == "down":
|
||||
return "looking strongly right and down"
|
||||
|
||||
return "looking forward"
|
||||
|
||||
|
||||
def _detect_objects(pil_img: Image.Image) -> list:
|
||||
def _estimate_bbox_for_tag(tag: str, keypoints: list, width: int, height: int, alpha_bbox: list = None) -> list:
|
||||
"""Estimate a bounding box for a given tag using pose keypoints or alpha bounding box.
|
||||
|
||||
Returns a list [x1, y1, x2, y2] of pixel coordinates, or None.
|
||||
"""
|
||||
import math
|
||||
if not keypoints or len(keypoints) < 17:
|
||||
if alpha_bbox:
|
||||
return [int(v) for v in alpha_bbox]
|
||||
return [0, 0, width, height]
|
||||
|
||||
tag_lower = tag.lower().replace("_", " ")
|
||||
|
||||
# Define terms lists mapped to anatomical structures
|
||||
head_terms = ["hair", "head", "face", "eye", "eyes", "nose", "ear", "ears", "mouth", "makeup", "eyebrow", "eyebrows", "glasses", "sunglasses", "earrings", "jewelry", "blush", "necklace", "collar", "hat", "cap", "crown", "smile", "gaze", "cheek", "teeth", "lips"]
|
||||
chest_terms = ["breast", "nipple", "nipples", "breasts", "chest", "cleavage", "bra", "bikini top", "top", "shirt", "collarbone", "pendant"]
|
||||
stomach_terms = ["navel", "stomach", "belly", "abs", "midriff", "waist", "panties", "underwear", "pelvis", "bikini bottom", "hips", "hip"]
|
||||
arm_terms = ["arm", "arms", "hand", "hands", "wrist", "wrists", "elbow", "elbows", "finger", "fingers", "sleeve", "sleeves", "glove", "gloves"]
|
||||
leg_terms = ["leg", "legs", "thigh", "thighs", "knee", "knees", "calf", "calves", "foot", "feet", "ankle", "ankles", "shoe", "shoes", "socks", "sock", "boots", "boot"]
|
||||
|
||||
bbox = None
|
||||
|
||||
# 1. Head/Face
|
||||
if any(term in tag_lower for term in head_terms):
|
||||
head_kpts = [keypoints[i] for i in range(5) if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||||
if head_kpts:
|
||||
xs = [kp[0] for kp in head_kpts]
|
||||
ys = [kp[1] for kp in head_kpts]
|
||||
min_x, max_x = min(xs), max(xs)
|
||||
min_y, max_y = min(ys), max(ys)
|
||||
|
||||
# Determine padding based on shoulder distance if available
|
||||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||||
shoulder_dist = math.dist(keypoints[5][:2], keypoints[6][:2])
|
||||
pad_x = shoulder_dist * 0.35
|
||||
pad_y = shoulder_dist * 0.45
|
||||
else:
|
||||
pad_x = max(max_x - min_x, width * 0.08)
|
||||
pad_y = max(max_y - min_y, height * 0.08)
|
||||
|
||||
bbox = [
|
||||
min_x - pad_x,
|
||||
min_y - pad_y * 1.3, # pull top higher to cover hair/hats
|
||||
max_x + pad_x,
|
||||
max_y + pad_y * 0.7
|
||||
]
|
||||
|
||||
# 2. Chest/Breasts
|
||||
elif any(term in tag_lower for term in chest_terms):
|
||||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||||
sh_mid_x = (keypoints[5][0] + keypoints[6][0]) / 2
|
||||
sh_mid_y = (keypoints[5][1] + keypoints[6][1]) / 2
|
||||
sh_dist = math.dist(keypoints[5][:2], keypoints[6][:2])
|
||||
|
||||
if len(keypoints) > 12 and keypoints[11][2] >= 0.3 and keypoints[12][2] >= 0.3:
|
||||
hip_mid_y = (keypoints[11][1] + keypoints[12][1]) / 2
|
||||
torso_h = hip_mid_y - sh_mid_y
|
||||
else:
|
||||
torso_h = sh_dist * 1.2
|
||||
|
||||
chest_center_y = sh_mid_y + torso_h * 0.28
|
||||
chest_w = sh_dist * 0.95
|
||||
chest_h = torso_h * 0.42
|
||||
bbox = [
|
||||
sh_mid_x - chest_w / 2,
|
||||
chest_center_y - chest_h / 2,
|
||||
sh_mid_x + chest_w / 2,
|
||||
chest_center_y + chest_h / 2
|
||||
]
|
||||
|
||||
# 3. Midriff/Pelvis/Hips/Underwear
|
||||
elif any(term in tag_lower for term in stomach_terms):
|
||||
if len(keypoints) > 12 and keypoints[11][2] >= 0.3 and keypoints[12][2] >= 0.3:
|
||||
hip_mid_x = (keypoints[11][0] + keypoints[12][0]) / 2
|
||||
hip_mid_y = (keypoints[11][1] + keypoints[12][1]) / 2
|
||||
hip_dist = math.dist(keypoints[11][:2], keypoints[12][:2])
|
||||
|
||||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||||
sh_mid_y = (keypoints[5][1] + keypoints[6][1]) / 2
|
||||
torso_h = hip_mid_y - sh_mid_y
|
||||
else:
|
||||
torso_h = hip_dist * 1.5
|
||||
|
||||
if any(term in tag_lower for term in ["navel", "stomach", "belly", "abs", "midriff", "waist"]):
|
||||
center_y = hip_mid_y - torso_h * 0.22
|
||||
box_h = torso_h * 0.32
|
||||
box_w = hip_dist * 1.15
|
||||
else: # panties, underwear, pelvis, bikini bottom, hips
|
||||
center_y = hip_mid_y + torso_h * 0.05
|
||||
box_h = torso_h * 0.42
|
||||
box_w = hip_dist * 1.25
|
||||
|
||||
bbox = [
|
||||
hip_mid_x - box_w / 2,
|
||||
center_y - box_h / 2,
|
||||
hip_mid_x + box_w / 2,
|
||||
center_y + box_h / 2
|
||||
]
|
||||
|
||||
# 4. Arms
|
||||
elif any(term in tag_lower for term in arm_terms):
|
||||
arm_indices = [5, 6, 7, 8, 9, 10]
|
||||
xs = [keypoints[i][0] for i in arm_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||||
ys = [keypoints[i][1] for i in arm_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||||
if xs:
|
||||
min_x, max_x = min(xs), max(xs)
|
||||
min_y, max_y = min(ys), max(ys)
|
||||
pad_x = (max_x - min_x) * 0.12 + 15
|
||||
pad_y = (max_y - min_y) * 0.12 + 15
|
||||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||||
|
||||
# 5. Legs
|
||||
elif any(term in tag_lower for term in leg_terms):
|
||||
leg_indices = [11, 12, 13, 14, 15, 16]
|
||||
xs = [keypoints[i][0] for i in leg_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||||
ys = [keypoints[i][1] for i in leg_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||||
if xs:
|
||||
min_x, max_x = min(xs), max(xs)
|
||||
min_y, max_y = min(ys), max(ys)
|
||||
pad_x = (max_x - min_x) * 0.12 + 15
|
||||
pad_y = (max_y - min_y) * 0.12 + 15
|
||||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||||
|
||||
# 6. Fallback or general full-body
|
||||
if bbox is None:
|
||||
visible_kpts = [k for k in keypoints if k[2] >= 0.3]
|
||||
if visible_kpts:
|
||||
xs = [k[0] for k in visible_kpts]
|
||||
ys = [k[1] for k in visible_kpts]
|
||||
min_x, max_x = min(xs), max(xs)
|
||||
min_y, max_y = min(ys), max(ys)
|
||||
pad_x = (max_x - min_x) * 0.15 + 20
|
||||
pad_y = (max_y - min_y) * 0.12 + 20
|
||||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||||
elif alpha_bbox:
|
||||
bbox = list(alpha_bbox)
|
||||
else:
|
||||
bbox = [0, 0, width, height]
|
||||
|
||||
# Clip to image bounds and format
|
||||
x1 = max(0, min(int(bbox[0]), width))
|
||||
y1 = max(0, min(int(bbox[1]), height))
|
||||
x2 = max(0, min(int(bbox[2]), width))
|
||||
y2 = max(0, min(int(bbox[3]), height))
|
||||
|
||||
# Ensure some valid size
|
||||
if x2 <= x1:
|
||||
x2 = min(x1 + 10, width)
|
||||
if y2 <= y1:
|
||||
y2 = min(y1 + 10, height)
|
||||
|
||||
return [x1, y1, x2, y2]
|
||||
|
||||
|
||||
def _detect_objects(pil_img: Image.Image, keypoints: list = None) -> list:
|
||||
"""Detect objects in the image using WD tagger.
|
||||
|
||||
Returns a list of detected objects with bounding box coordinates.
|
||||
@@ -6180,14 +6568,20 @@ def _detect_objects(pil_img: Image.Image) -> list:
|
||||
|
||||
# Filter for object-related tags (general and character categories)
|
||||
objects = []
|
||||
|
||||
width, height = pil_img.size
|
||||
alpha_bbox = None
|
||||
if pil_img.mode == 'RGBA':
|
||||
alpha = pil_img.split()[-1]
|
||||
alpha_bbox = alpha.getbbox()
|
||||
|
||||
for t in tags:
|
||||
if t["cat"] in (0, 4): # general and character categories
|
||||
# For simplicity, we'll return just the tag name with confidence
|
||||
# In a more advanced implementation, we could extract bounding boxes from the model
|
||||
bbox = _estimate_bbox_for_tag(t["tag"], keypoints, width, height, alpha_bbox)
|
||||
objects.append({
|
||||
"tag": t["tag"],
|
||||
"score": t["score"],
|
||||
"bbox": None # No bounding box available from WD tagger
|
||||
"bbox": bbox
|
||||
})
|
||||
return objects
|
||||
except Exception as e:
|
||||
@@ -6226,8 +6620,9 @@ def _process_image_for_metadata(filename: str):
|
||||
best_person = _best_person(people)
|
||||
|
||||
# Extract metadata
|
||||
width, height = pil_img.size
|
||||
people_count = _detect_people_count(best_person)
|
||||
anatomical_completeness = _detect_anatomical_completeness(best_person)
|
||||
anatomical_completeness = _detect_anatomical_completeness(best_person, width, height)
|
||||
facial_direction = _detect_facial_direction(best_person)
|
||||
|
||||
# ALSO extract pose description and pose skeleton
|
||||
@@ -6244,7 +6639,7 @@ def _process_image_for_metadata(filename: str):
|
||||
print(f"[pose] index save failed for {filename}: {e}")
|
||||
|
||||
# Detect objects
|
||||
objects = _detect_objects(pil_img)
|
||||
objects = _detect_objects(pil_img, keypoints=best_person)
|
||||
|
||||
# Update database with new metadata
|
||||
database.upsert_person(
|
||||
@@ -6276,6 +6671,7 @@ def _process_image_for_metadata(filename: str):
|
||||
|
||||
class BackfillMetadataRequest(BaseModel):
|
||||
filenames: list[str] | None = None # If None, process all images in DB
|
||||
force: bool = False
|
||||
|
||||
|
||||
import asyncio
|
||||
@@ -6283,6 +6679,21 @@ from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
|
||||
_metadata_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="metadata")
|
||||
|
||||
|
||||
@app.post("/images/invalidate-metadata")
|
||||
def invalidate_metadata():
|
||||
"""Invalidate all metadata records by setting their columns to NULL.
|
||||
This enables full reprocessing via backfill or the background idle loop.
|
||||
"""
|
||||
try:
|
||||
database.invalidate_all_metadata()
|
||||
_failed_backfill_filenames.clear()
|
||||
_invalidate_static()
|
||||
return {"status": "success", "message": "All earlier metadata has been invalidated successfully."}
|
||||
except Exception as e:
|
||||
print(f"[metadata] Invalidation error: {e}")
|
||||
raise HTTPException(500, f"Invalidation failed: {str(e)}")
|
||||
|
||||
|
||||
@app.post("/images/backfill-metadata")
|
||||
async def backfill_metadata(req: BackfillMetadataRequest):
|
||||
"""Backfill metadata for existing images in the database.
|
||||
@@ -6291,6 +6702,9 @@ async def backfill_metadata(req: BackfillMetadataRequest):
|
||||
people count, anatomical completeness, facial direction, and objects.
|
||||
"""
|
||||
try:
|
||||
if req.force:
|
||||
_failed_backfill_filenames.clear()
|
||||
|
||||
# Get list of all image files
|
||||
if req.filenames is not None:
|
||||
filenames = req.filenames
|
||||
|
||||
Reference in New Issue
Block a user