StickerGen / face_to_prompt.py
KarthikAI's picture
Update face_to_prompt.py
6022861 verified
import os
os.environ["INSIGHTFACE_HOME"] = "/data/.insightface"
os.makedirs("/data/.insightface", exist_ok=True)
os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
os.makedirs("/tmp/matplotlib", exist_ok=True)
print("INSIGHTFACE_HOME set to:", os.environ.get("INSIGHTFACE_HOME"))
from deepface import DeepFace
from PIL import Image
import insightface
import numpy as np
def extract_face_prompt(img_path):
# Load image as numpy array
img = Image.open(img_path).convert("RGB")
img_np = np.array(img)
# DeepFace analysis
deepface_result = DeepFace.analyze(
img_path, actions=['age', 'gender', 'race', 'emotion'], enforce_detection=False
)
if isinstance(deepface_result, list):
result = deepface_result[0]
else:
result = deepface_result
age = result['age']
gender = result['dominant_gender']
race = result['dominant_race']
emotion = result['dominant_emotion']
# InsightFace analysis
# model = insightface.app.FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
model = insightface.app.FaceAnalysis(
name='buffalo_l',
providers=['CPUExecutionProvider'],
root=os.environ.get("INSIGHTFACE_HOME", "/data/.insightface")
)
model.prepare(ctx_id=0)
faces = model.get(img_np)
face_attrs = {}
if faces:
face = faces[0]
face_attrs['glasses'] = face.get('glass', 0) > 0
face_attrs['beard'] = face.get('beard', 0) > 0
face_attrs['mustache'] = face.get('mustache', 0) > 0
attr_parts = []
if face_attrs.get('glasses'):
attr_parts.append("wearing glasses")
if face_attrs.get('beard'):
attr_parts.append("has a beard")
if face_attrs.get('mustache'):
attr_parts.append("has a mustache")
attr_str = ", ".join(attr_parts) if attr_parts else ""
prompt = (
f"A chibi-style digital sticker of a {gender.lower()} with {race} skin, "
f"{'smiling' if emotion=='happy' else emotion}, approximately {age} years old"
)
if attr_str:
prompt += f", {attr_str}"
prompt += ", closeup, cute, expressive face, white outline, transparent background"
return prompt